Enigma 密码破解

class BaseRotor:
    """定义转子基础类，存储转子的转换规则"""

    def __init__(self, wiring: str, notch: str):
        """初始化函数，传入转子的转换规则和转子的进位卡口位置"""
        self.wiring = wiring
        self.notch = ord(notch) - 65

# 按照题目要求，定义转子数据
ROTOR_1 = BaseRotor("EKMFLGDQVZNTOWYHXUSPAIBRCJ", "Q")
ROTOR_2 = BaseRotor("AJDKSIRUXBLHWTMCQGZNPYFVOE", "E")
ROTOR_3 = BaseRotor("BDFHJLCPRTXVZNYEIWGAKMUSQO", "V")
# 定义反射器数据
REFLECTOR = "YRUHQSLDPXNGOKMIEBFZCWVJAT"

def get_rotors(r1, r2, r3):
    """封装函数，通过转子的序号获得对应的转子"""
    rotors = [ROTOR_1, ROTOR_2, ROTOR_3]
    return [rotors[r1 - 1], rotors[r2 - 1], rotors[r3 - 1]]

class Rotor:
    """定义转子类，用于实际执行转子对字符的转换"""

    def __init__(self, rotor: BaseRotor, ring_setting: str, position: str):
        """初始化函数，传入转子的转换规则，以及转子的 Ring setting 和初始位置"""
        self.wiring = rotor.wiring
        self.notch = rotor.notch
        self.ring_setting = ord(ring_setting) - 65
        self.position = ord(position) - 65

    def forward(self, num: int) -> int:
        """转子的前向转换函数"""
        delta = self.position - self.ring_setting
        num = (num + delta) % 26
        num = ord(self.wiring[num]) - 65
        num = (num - delta) % 26
        return num

    def backward(self, num: int) -> int:
        """转子的后向转换函数"""
        delta = self.position - self.ring_setting
        num = (num + delta) % 26
        num = self.wiring.index(chr(num + 65))
        num = (num - delta) % 26
        return num

    def rotate(self) -> bool:
        """转子的进位函数"""
        flag = self.position == self.notch
        self.position = (self.position + 1) % 26
        return flag

class Enigma:
    """定义 Enigma 类，用于实际执行加密/解密"""

    def __init__(self, reflector: str, rotors: list[Rotor], rings: str, positions: str):
        """初始化函数，传入反射器的转换规则，以及三个转子的转换规则、Ring setting 和初始位置"""
        self.rotors = [Rotor(rotors[i], rings[i], positions[i])
                       for i in range(3)]
        self.reflector = reflector

    def rotate(self):
        """模拟转子的旋转和进位"""
        # 右侧转子旋转
        if self.rotors[2].rotate():
            # 如果右侧转子进位，中间转子旋转
            if self.rotors[1].rotate():
                # 如果中间转子进位，左侧转子旋转
                self.rotors[0].rotate()
        elif self.rotors[1].position == self.rotors[1].notch:
            # 机器存在 Double step 问题，模拟这一行为
            if self.rotors[1].rotate():
                # 如果中间转子进位，左侧转子旋转
                self.rotors[0].rotate()

    def encrypt_single(self, char: str):
        """加密单个字符"""
        # 将明文字符转换为数字
        num = ord(char) - 65
        # 从右至左依次经过三个转子
        for i in range(2, -1, -1):
            num = self.rotors[i].forward(num)
        # 通过反射器进行转换
        num = ord(self.reflector[num]) - 65
        # 从左至右依次经过三个转子
        for i in range(3):
            num = self.rotors[i].backward(num)
        # 将数字转换为密文字符
        return chr(num + 65)

    def encrypt(self, plaintext: str):
        """加密字符串"""
        buffer = ""
        for char in plaintext:
            self.rotate()
            buffer += self.encrypt_single(char)

        return buffer

    def generate_chain(self):
        """对某一机器状态，生成A-Z对应的密文"""
        buffer = ""
        for i in range(26):
            char = chr(i + 65)
            buffer += self.encrypt_single(char)
        return buffer

def init_positions():
    """生成所有可能的机器状态"""
    for i in range(26):
        for j in range(26):
            for k in range(26):
                yield chr(i + 65) + chr(j + 65) + chr(k + 65)

if __name__ == "__main__":
    # 检查一下 Enigma 机和题目里的是不是一样的。这里是没有接线板的结果。
    enigma = Enigma(REFLECTOR, get_rotors(2, 3, 1), "DES", "AAA")
    assert enigma.encrypt("ASTWBKDJXSAXQFCJT") == "NVYEAYJBLRPBIRJOD"

    print("请运行 rejewski.py 或 turing.py")

# 使用标准库 itertools 辅助函数
from itertools import permutations
from enigma import Enigma, REFLECTOR, get_rotors, init_positions

def find_chain(alphabet, letters):
    """按照 Rejewski 的策略，在两行密文中找到字母链条"""
    groups = []
    grouped = alphabet[0]
    current_group = alphabet[0]

    while len(grouped) < 26:
        letter = current_group[-1]
        dst = letters[alphabet.index(letter)]
        if dst in current_group:
            groups.append(current_group)
            current_group = ""
            for i in range(26):
                if not alphabet[i] in grouped:
                    current_group = alphabet[i]
                    grouped += current_group
                    break
        else:
            grouped += dst
            current_group += dst
    if current_group:
        groups.append(current_group)
    return groups

def chain_fingerprint(rows):
    """将字母链条的数量和长度，作为机器状态的指纹"""
    chains = []
    for i in range(3):
        chains.append(find_chain(rows[i], rows[i + 3]))
    return [sorted(map(lambda s: len(s), chain)) for chain in chains]

def enigma_fingerprint(enigma):
    """对特定的机器生成指纹"""
    rows = []
    for _ in range(6):
        enigma.rotate()
        rows.append(enigma.generate_chain())
    return chain_fingerprint(rows)

def rejewski_crack():
    """按照 Rejewski 的策略进行破解"""
    # 题目中给出的第 4 5 6 行密文
    configs = [
        "ELCONWDIAPKSZHFBQTJYRGVXMU",
        "MRWJFDVSQEXUCONHBIPLTGAYZK",
        "WADFRPOLNTVCHMYBJQIGEUSKZX"
    ]
    # 第 1 2 3 行密文
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    expected = chain_fingerprint([alphabet] * 3 + configs)
    # 固定 Ring setting，恢复转子顺序和 Initial position
    for r in permutations(range(1, 4)):
        print("Trying rotor order", r)
        rotors = get_rotors(*r)
        for p in init_positions():
            enigma = Enigma(
                REFLECTOR, rotors, "DES", p)
            # 指纹符合，则认为找到了正确的机器状态
            if enigma_fingerprint(enigma) == expected:
                print("Solution found:", r, p)

if __name__ == "__main__":
    rejewski_crack()

# 使用 networkx 库，寻找图中的环路
import networkx as nx
# 使用标准库 itertools 辅助函数
from itertools import permutations, product
from enigma import Enigma, REFLECTOR, get_rotors, init_positions

class Cycle():
    """封装一个类，用于表示 Turing 的已知明文攻击的方法中找到的圈"""

    def __init__(self, letters) -> None:
        """初始化函数，记录圈中的字母"""
        self.letters = letters
        self.offsets = []

    def __repr__(self):
        return f"Cycle({self.letters})"

def find_cycle():
    """由题目中给出的已知明文和密文，找到其中所有的圈"""
    a = "ASTWBKDJXSAXQFCJT"
    b = "STWTJDJXSAQBJQJOD"
    edges = [(a[i], b[i]) for i in range(len(a))]
    g = nx.DiGraph(edges)

    cycles = []
    for c in nx.simple_cycles(g):
        cycle = Cycle(c)
        for i in range(len(c)):
            cycle.offsets.append(edges.index((c[i], c[(i + 1) % len(c)])))
        cycles.append(cycle)
    cycles.sort(key=lambda c: len(c.letters))
    return cycles

def get_plugboard_constraint(cycle, guess):
    """
    虽然题目不要求恢复接线板，但在猜测字母圈的过程中，自然会得到接线板存在的规则
    它们可以用作约束条件，判断猜测的转轮初值是否正确
    """
    mapping = {}
    for index, l in enumerate(cycle.letters):
        if l in mapping:
            if not mapping[l] == guess[index]:
                return None
        elif guess[index] in mapping:
            if not mapping[guess[index]] == l:
                return None
        else:
            mapping[l] = guess[index]
            mapping[guess[index]] = l

    return mapping

def examine_constraints(constraints):
    """
    检查我们猜测的接线板规则是否存在冲突
    例如，为了满足第一个圈，我们需要假定 A 连接到了 B
    而为了满足第二个圈，我们需要假定 A 连接到了 C
    这两个规则是冲突的，因此我们需要放弃这个猜测
    """
    for constraint in product(*constraints):
        mapping = {}
        succeeded = True
        for c in constraint:
            for k, v in c.items():
                if k in mapping:
                    if not mapping[k] == v:
                        succeeded = False
                elif v in mapping:
                    if not mapping[v] == k:
                        succeeded = False
                else:
                    mapping[k] = v
        if succeeded:
            return True

def apply_cycles(mappings, cycles):
    """检查某种转轮初始值是否满足全部的字母圈"""
    constraints = []
    for cycle in cycles:
        constraint = []
        for i in range(26):
            guess = [chr(i + 65)]
            for j in cycle.offsets:
                guess.append(mappings[j][ord(guess[-1]) - 65])
            if guess[-1] == guess[0]:
                c = get_plugboard_constraint(cycle, guess)
                if c:
                    constraint.append(c)
        if not len(constraint):
            return False
        else:
            constraints.append(constraint)
    return examine_constraints(constraints)

def turing_crack():
    """按照 Turing 的策略进行破解"""
    # 固定 Ring setting，恢复转子顺序和 Initial position
    for r in permutations(range(1, 4)):
        print("Trying rotor order", r)
        rotors = get_rotors(*r)
        for p in init_positions():
            enigma = Enigma(
                REFLECTOR, rotors, "DES", p)
            mappings = []

            cycles = find_cycle()
            for _ in range(max(map(lambda c: max(c.offsets), cycles)) + 1):
                enigma.rotate()
                mappings.append(enigma.generate_chain())

            if apply_cycles(mappings, cycles):
                print("Solution found:", r, p)

if __name__ == "__main__":
    turing_crack()