Solution AI-CTF-2022 task RNN Generator

decryptor.py

from itertools import product
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import torch
from tqdm import tqdm

from encryptor import Encryptor
from generator import _EXPECTED_VALUES, REGISTER_SIZE, S0, S1, Register


def get_bit_array_from_bytes_array(array):
    return torch.tensor([list(map(int, "{0:08b}".format(i))) for i in array]).flatten()


class Mitm:
    """
    Meet in the middle realisation class
    """

    def __init__(self, known_gamma_length: int):
        self._founded_keys: List[Tuple[torch.Tensor, torch.Tensor]] = []
        self.known_gamma_length: int = known_gamma_length
        self._gamma_dict: Dict[Tuple[int]] = dict()
        self._init_gamma_dict()

    @staticmethod
    def _all_init_states() -> Iterable:
        """
        Generate all keys(init states for register)
        :return:
        """
        return product(*[_EXPECTED_VALUES for _ in range(REGISTER_SIZE)])

    def _init_gamma_dict(self) -> None:
        """
        Fill dictionary for mitm method
        :return:
        """
        for x_0_init in tqdm(self._all_init_states(), total=2**REGISTER_SIZE):
            register_0 = Register(torch.FloatTensor([x_0_init]), S0)
            f_0_segment = tuple(
                register_0.generate_bit() for _ in range(self.known_gamma_length)
            )

            self._gamma_dict[f_0_segment] = x_0_init

    def find_key(self, gamma_segment: torch.Tensor) -> None:
        """
        Find all potential keys with mitm method
        :param gamma_segment: array with gamma segment
        :return:
        """
        self._founded_keys = []
        for x_1_init in tqdm(self._all_init_states(), total=2**REGISTER_SIZE):

            register_1 = Register(torch.FloatTensor([x_1_init]), S1)
            f_1_segment = torch.Tensor(
                [register_1.generate_bit() for _ in range(self.known_gamma_length)]
            )

            f_0_segment = tuple(((gamma_segment + f_1_segment) % 2).tolist())

            if f_0_segment in self._gamma_dict:
                print("FIND KEY!!!")
                self._founded_keys.append(
                    (
                        torch.Tensor([self._gamma_dict[f_0_segment]]),
                        torch.Tensor([x_1_init]),
                    )
                )

    @property
    def keys(self) -> List[Tuple[torch.Tensor, torch.Tensor]]:
        """
        Keys found with mitm method
        :return:
        """
        return self._founded_keys

    def decrypt_ogg_file(
        self, encrypted_file_path: Path, decrypted_files_dir: Path
    ) -> None:
        """
        Decrypt ogg sound file
        :param encrypted_file_path: cypher text file path
        :param decrypted_files_dir: directory with results, may be not one
        :return:
        """
        ogg_header = [79, 103, 103, 83]

        with open(encrypted_file_path, "rb") as f:
            cypher_text = f.read()

        gamma_bytes = []
        for ot, ct in zip(ogg_header[:4], cypher_text[:4]):
            gamma_bytes.append((int(ct) + ot) % 256)

        gamma_bit_segment = get_bit_array_from_bytes_array(gamma_bytes)

        self.find_key(gamma_segment=gamma_bit_segment)

        for i, (x_0_init, x_1_init) in enumerate(self._founded_keys):
            enc = Encryptor(x_0_init, x_1_init)
            enc.apply(encrypted_file_path, Path(f"{decrypted_files_dir}/{i}.ogg"))