罚罪相当。罚，punishment; 罪， crime. 相当means equal. It is fit to your meaning , but I am not sure if it is a chengyu, The chengyu “咎由自取”or “罪有应得”， both means that a man is deserved for his crime, but not to mention what kind of crime. There are also a chengyu, 以毁为罚yǐ huǐ wéi fá, mean that the punishment should be measured according to his crime.

It is a question in make a decision whether the crime fit into the definition of law. Argument: Some of these crimes weren't crimes in the past and some activities which used to be crimes are legal today. There was a time when you couldn't drink alcohol but you could legally smoke marijuana. In an even earlier era, you could own slaves but if you were a woman you couldn't own real estate. 'the punishment should fit the crime" interpreted as the crimes that fit into the definition of today's law.

罪有应得

可以說 罪刑相當 罪刑相稱 罰當其罪 或者 依罪量刑 依罪科罰 依罪量刑 可能比較接近你的期望

This site will help you in future; Dictionary of more than 30 000 Chengyu (chinese proverbial idioms) with calligraphy, pinyin, english translation, chinese explanations, synonyms, ... www.chinese-tools.com/chinese/chengyu/dictionary