test_cipher_consistency.py

import unittest

import numpy as np

import step1_generate_data as step1
import step2_run_ml as step2
import step3_key_recovery as step3


class LastRoundConsistencyTests(unittest.TestCase):
    def test_round_parameters_are_explicit_and_greedy_pool_is_sufficient(self):
        self.assertIn(11, step2.ROUND_PARAMS)
        for rounds, (n_target, *_rest) in step2.ROUND_PARAMS.items():
            self.assertLess(n_target, step2.GREEDY_POOL, msg=f"round {rounds}")
        with self.assertRaises(ValueError):
            step2.get_round_params(12)

    def test_partial_decryption_matches_ground_truth_pre_final_sbox_state(self):
        rng = np.random.default_rng(1234)
        pbox_inv = [0] * 64
        for src, dst in enumerate(step1.PBOX):
            pbox_inv[dst] = src
        sbox_inv = [step1.SBOX.index(v) for v in range(16)]

        for _ in range(100):
            key = step1.rand80(rng)
            round_keys = step1.key_schedule_80(key)
            p1 = step1.rand64(rng)
            p2 = p1 ^ 0x0000000000800000
            c1 = step1.present_encrypt(p1, round_keys, 6)
            c2 = step1.present_encrypt(p2, round_keys, 6)

            def pre_final_sbox_state(plaintext):
                state = plaintext
                for r in range(5):
                    state = step1.pbox_layer(step1.sbox_layer(state ^ round_keys[r]))
                return state ^ round_keys[5]

            expected1 = pre_final_sbox_state(p1)
            expected2 = pre_final_sbox_state(p2)
            self.assertEqual(
                step2.partial_decrypt_last(c1, round_keys[6], sbox_inv, pbox_inv),
                expected1,
            )
            self.assertEqual(
                step2.partial_decrypt_last(c2, round_keys[6], sbox_inv, pbox_inv),
                expected2,
            )
            actual = step3._compute_diffs(
                np.asarray([round_keys[6]], dtype=np.uint64),
                np.asarray([c1], dtype=np.uint64),
                np.asarray([c2], dtype=np.uint64),
            )[0, 0]
            self.assertEqual(int(actual), expected1 ^ expected2)


if __name__ == "__main__":
    unittest.main()