porta_cracker.py

from matplotlib import pyplot as plt
from tabulate import tabulate

from chi_squared import ChiSquared
from stats_collector import StatsCollector

tableau = {
    0: "NOPQRSTUVWXYZABCDEFGHIJKLM",
    1: "OPQRSTUVWXYZNMABCDEFGHIJKL",
    2: "PQRSTUVWXYZNOLMABCDEFGHIJK",
    3: "QRSTUVWXYZNOPKLMABCDEFGHIJ",
    4: "RSTUVWXYZNOPQJKLMABCDEFGHI",
    5: "STUVWXYZNOPQRIJKLMABCDEFGH",
    6: "TUVWXYZNOPQRSHIJKLMABCDEFG",
    7: "UVWXYZNOPQRSTGHIJKLMABCDEF",
    8: "VWXYZNOPQRSTUFGHIJKLMABCDE",
    9: "WXYZNOPQRSTUVEFGHIJKLMABCD",
    10: "XYZNOPQRSTUVWDEFGHIJKLMABC",
    11: "YZNOPQRSTUVWXCDEFGHIJKLMAB",
    12: "ZNOPQRSTUVWXYBCDEFGHIJKLMA"
}


class PortaCracker:
    def __init__(self, cipher_statistics: StatsCollector, key_length, label="Porta Tableau Errors"):
        self.cipher_statistics = cipher_statistics
        self.key_length = key_length
        self.label = label

    @staticmethod
    def encrypt(key, plain_text):
        key = key.upper()
        plain_text = plain_text.upper()
        return ''.join(tableau[(ord(key[i % len(key)]) - ord('A')) // 2][ord(c) - ord('A')]
                       for i, c in enumerate(plain_text))

    @staticmethod
    def decrypt(key, cipher_text):
        return PortaCracker.encrypt(key, cipher_text)

    def crack(self):
        substrings = self._substrings()
        key = []
        for index_s, s in enumerate(substrings):
            print(f"Attempt for key element {index_s}")
            tableau_index_of_key = self._crack_key_element(s, index_s)
            key_element = PortaCracker._tableau_index_to_key_element(tableau_index_of_key)
            print(f"Found: {key_element}")
            key.append(key_element)
        return key

    def _crack_key_element(self, s, key_element_id=0):
        stat = []
        for i in range(len(tableau)):
            decrypted_s = self.decrypt(chr(i * 2 + ord("A")), s)
            chi_squared = ChiSquared(decrypted_s).val
            stat.append(chi_squared)
        PortaCracker._pretty_print(stat)
        PortaCracker.plot_key_element_error(stat, key_element_id)
        return stat.index(min(stat))

    def _substrings(self):
        return [
            "".join(
                self.cipher_statistics.data[j]
                for j in range(i, self.cipher_statistics.total_chars, self.key_length)
            )
            for i in range(self.key_length)
        ]

    @staticmethod
    def _pretty_print(stat):
        pretty_stat = [["Key Element", "Chi Squared"]]
        pretty_stat.extend([PortaCracker._tableau_index_to_key_element(i), stat[i]] for i in range(len(stat)))
        print(tabulate(pretty_stat, headers="firstrow", tablefmt="rounded_grid"))

    @staticmethod
    def _tableau_index_to_key_element(i):
        return chr(i * 2 + ord("A")) + "-" + chr(i * 2 + ord("A") + 1)

    @staticmethod
    def plot_key_element_error(stat, key_element_id):
        min_chi_squared = min(stat)
        pretty_stat = {
            PortaCracker._tableau_index_to_key_element(i): stat[i]
            for i in range(len(stat))
        }
        plt.plot(pretty_stat.keys(), pretty_stat.values())
        plt.xlabel(
            f"Key Element {key_element_id} Found: {PortaCracker._tableau_index_to_key_element(stat.index(min(stat)))}")
        plt.ylabel("Chi Squared")
        plt.title(f"Porta Tableau Errors for Key Element: {key_element_id}")

        plt.axhline(y=min_chi_squared, color='r', linestyle='--',
                    label=f'Min Chi Squared ({min_chi_squared})')

        plt.legend()
        plt.show()