TMC #0004: Enigma Machine and Bombe Attack Simulator

"""
TMC #0004 — Turing's Bombe and the Breaking of Enigma
======================================================
Simulates the core cryptanalytic machinery of:
  1. The Enigma machine (3-rotor Wehrmacht variant)
  2. Turing's Bombe — the electro-mechanical device
     that exploited Enigma's structural weaknesses
  3. The crib-dragging attack (known-plaintext search)

Historical note: The real Bombe was electro-mechanical,
running at ~120 RPM, testing 1 rotor position per cycle.
This Python model captures the *algorithm* faithfully.

Run:
    python bombe.py
"""

from __future__ import annotations

import itertools
import string
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

ALPHA = string.ascii_uppercase  # 26-letter alphabet
N = 26


# ─────────────────────────────────────────────
# 1. Enigma Components
# ─────────────────────────────────────────────


def make_wiring(mapping: str) -> List[int]:
    """Convert alphabet string to forward permutation list."""
    return [ALPHA.index(c) for c in mapping.upper()]


def invert_wiring(wiring: List[int]) -> List[int]:
    """Invert a permutation."""
    inv = [0] * N
    for i, w in enumerate(wiring):
        inv[w] = i
    return inv


# Historical Enigma rotor wirings (Wehrmacht/Luftwaffe)
ROTOR_WIRINGS = {
    "I": make_wiring("EKMFLGDQVZNTOWYHXUSPAIBRCJ"),
    "II": make_wiring("AJDKSIRUXBLHWTMCQGZNPYFVOE"),
    "III": make_wiring("BDFHJLCPRTXVZNYEIWGAKMUSQO"),
    "IV": make_wiring("ESOVPZJAYQUIRHXLNFTGKDCMWB"),
    "V": make_wiring("VZBRGITYUPSDNHLXAWMJQOFECK"),
}

ROTOR_NOTCHES = {
    "I": 16,  # Q → R turnover
    "II": 4,  # E → F turnover
    "III": 21,  # V → W turnover
    "IV": 9,  # J → K turnover
    "V": 25,  # Z → A turnover
}

REFLECTOR_B = make_wiring("YRUHQSLDPXNGOKMIEBFZCWVJAT")


@dataclass
class Rotor:
    name: str
    wiring: List[int]
    notch: int
    position: int = 0
    ring_setting: int = 0

    @property
    def at_notch(self) -> bool:
        return self.position == self.notch

    def step(self) -> None:
        self.position = (self.position + 1) % N

    def encode_forward(self, c: int) -> int:
        offset = (c + self.position - self.ring_setting) % N
        wired = self.wiring[offset]
        return (wired - self.position + self.ring_setting) % N

    def encode_backward(self, c: int) -> int:
        inv = invert_wiring(self.wiring)
        offset = (c + self.position - self.ring_setting) % N
        wired = inv[offset]
        return (wired - self.position + self.ring_setting) % N


def make_rotor(name: str, position: int = 0, ring: int = 0) -> Rotor:
    return Rotor(
        name=name,
        wiring=ROTOR_WIRINGS[name],
        notch=ROTOR_NOTCHES[name],
        position=position,
        ring_setting=ring,
    )


@dataclass
class Plugboard:
    """Steckerbrett: pairs of letters swapped before and after rotors."""

    swaps: Dict[int, int] = field(default_factory=dict)

    @staticmethod
    def from_pairs(pairs: List[str]) -> "Plugboard":
        pb = Plugboard()
        for pair in pairs:
            a, b = ALPHA.index(pair[0].upper()), ALPHA.index(pair[1].upper())
            pb.swaps[a] = b
            pb.swaps[b] = a
        return pb

    def encode(self, c: int) -> int:
        return self.swaps.get(c, c)


class EnigmaMachine:
    """
    Three-rotor Enigma machine (Wehrmacht variant).
    Rotors listed left-to-right (slow → fast).
    """

    def __init__(
        self,
        rotors: List[Rotor],
        reflector: List[int],
        plugboard: Plugboard,
    ):
        assert len(rotors) == 3
        self.rotors = rotors  # [left, middle, right]
        self.reflector = reflector
        self.plugboard = plugboard

    def _step_rotors(self) -> None:
        """Double-stepping anomaly included."""
        r, m, l = self.rotors[2], self.rotors[1], self.rotors[0]
        if m.at_notch:
            m.step()
            l.step()
        elif r.at_notch:
            m.step()
        r.step()

    def encrypt_char(self, c: int) -> int:
        """Encrypt one character index (0–25). Steps rotors first."""
        self._step_rotors()
        c = self.plugboard.encode(c)
        for rotor in reversed(self.rotors):  # right → left
            c = rotor.encode_forward(c)
        c = self.reflector[c]
        for rotor in self.rotors:  # left → right
            c = rotor.encode_backward(c)
        c = self.plugboard.encode(c)
        return c

    def encrypt(self, plaintext: str) -> str:
        result = []
        for ch in plaintext.upper():
            if ch in ALPHA:
                result.append(ALPHA[self.encrypt_char(ALPHA.index(ch))])
        return "".join(result)

    def set_positions(self, positions: Tuple[int, int, int]) -> None:
        for rotor, pos in zip(self.rotors, positions):
            rotor.position = pos


def make_enigma(
    rotor_names: Tuple[str, str, str] = ("I", "II", "III"),
    positions: Tuple[int, int, int] = (0, 0, 0),
    ring_settings: Tuple[int, int, int] = (0, 0, 0),
    plugboard_pairs: Optional[List[str]] = None,
) -> EnigmaMachine:
    rotors = [
        make_rotor(rotor_names[i], positions[i], ring_settings[i]) for i in range(3)
    ]
    pb = Plugboard.from_pairs(plugboard_pairs or [])
    return EnigmaMachine(rotors, REFLECTOR_B, pb)


# ─────────────────────────────────────────────
# 2. Bombe Core Principle
#
# The Bombe exploited one fatal flaw:
# Enigma NEVER encrypts a letter to itself.
# (The reflector guaranteed this.)
#
# Given a crib (known plaintext fragment),
# we can overlay it against the ciphertext
# at each offset and check for contradictions.
# Any position where a letter maps to itself
# is IMPOSSIBLE → eliminated.
#
# This is "crib dragging."
# ─────────────────────────────────────────────


def crib_drag(ciphertext: str, crib: str) -> List[int]:
    """
    Return all starting positions where the crib could plausibly
    encrypt to the corresponding ciphertext fragment —
    i.e., no position has the same letter in both crib and cipher.
    """
    ct = ciphertext.upper()
    cb = crib.upper()
    possible = []
    for start in range(len(ct) - len(cb) + 1):
        fragment = ct[start : start + len(cb)]
        # Enigma property: cipher[i] != plain[i] always
        if all(fragment[i] != cb[i] for i in range(len(cb))):
            possible.append(start)
    return possible


# ─────────────────────────────────────────────
# 3. Simplified Bombe Attack
#
# For small key spaces, brute-force the rotor
# start positions and check if decryption of the
# ciphertext contains the crib at the expected
# offset. Real Bombes tested menu-derived loops;
# this models the outer search loop faithfully.
# ─────────────────────────────────────────────


def bombe_attack(
    ciphertext: str,
    crib: str,
    crib_offset: int,
    rotor_config: Tuple[str, str, str],
    plugboard_pairs: Optional[List[str]] = None,
    verbose: bool = False,
) -> Optional[Tuple[int, int, int]]:
    """
    Search all 26^3 = 17,576 starting positions for a 3-rotor Enigma.
    Returns the first (left, middle, right) start position where
    decrypt(ciphertext)[crib_offset:] starts with crib.
    """
    plugboard_pairs = plugboard_pairs or []
    attempts = 0

    for l, m, r in itertools.product(range(N), repeat=3):
        attempts += 1
        machine = make_enigma(rotor_config, (l, m, r), plugboard_pairs=plugboard_pairs)
        decrypted = machine.encrypt(ciphertext)  # Enigma is symmetric

        candidate = decrypted[crib_offset : crib_offset + len(crib)]
        if candidate == crib.upper():
            if verbose:
                print(
                    f"  Hit at position ({ALPHA[l]}{ALPHA[m]}{ALPHA[r]}) "
                    f"after {attempts} attempts"
                )
                print(f"  Decrypted: {decrypted}")
            return (l, m, r)

    return None


# ─────────────────────────────────────────────
# 4. Demonstrations
# ─────────────────────────────────────────────


def demo_enigma_symmetry():
    print("=" * 60)
    print("Enigma Symmetry: encrypt(encrypt(x)) == x")
    print("=" * 60)
    plaintext = "HELLOTURING"
    m1 = make_enigma(positions=(4, 7, 19))
    ciphertext = m1.encrypt(plaintext)

    m2 = make_enigma(positions=(4, 7, 19))
    recovered = m2.encrypt(ciphertext)

    print(f"  Plaintext:  {plaintext}")
    print(f"  Ciphertext: {ciphertext}")
    print(f"  Recovered:  {recovered}")
    assert recovered == plaintext, "Symmetry failed!"
    print("  ✓ Symmetry confirmed\n")


def demo_no_self_encrypt():
    print("=" * 60)
    print("Enigma Fatal Flaw: a letter never encrypts to itself")
    print("=" * 60)
    machine = make_enigma()
    violations = 0
    for _ in range(100):
        for i in range(N):
            enc = machine.encrypt_char(i)
            if enc == i:
                violations += 1
    print(f"  Self-encryptions in 2600 trials: {violations}  (must be 0)")
    assert violations == 0
    print("  ✓ Confirmed — Enigma never maps a letter to itself\n")


def demo_crib_drag():
    print("=" * 60)
    print("Crib Dragging")
    print("=" * 60)
    # Bletchley Park knew German operators often started messages with
    # "WETTER" (weather) or "KEINE BESONDEREN EREIGNISSE" (nothing to report)
    machine = make_enigma(positions=(3, 14, 21))
    plaintext = "WETTERNICHTSBESONDERES"  # "weather nothing special"
    ciphertext = machine.encrypt(plaintext)

    print(f"  Plaintext:  {plaintext}")
    print(f"  Ciphertext: {ciphertext}")

    crib = "WETTER"
    possible = crib_drag(ciphertext, crib)
    print(f"  Crib '{crib}' possible offsets (no self-maps): {possible}")
    print()


def demo_bombe_attack():
    print("=" * 60)
    print("Bombe Attack: Recover start position from known crib")
    print("=" * 60)
    # Set up a known configuration (pretend attacker doesn't know position)
    true_positions = (5, 12, 3)  # E, M, D
    machine = make_enigma(("I", "II", "III"), true_positions)
    plaintext = "ANGRIFFMORGENFRUH"  # "attack tomorrow morning"
    ciphertext = machine.encrypt(plaintext)

    print(f"  True start positions: {tuple(ALPHA[p] for p in true_positions)}")
    print(f"  Plaintext:  {plaintext}")
    print(f"  Ciphertext: {ciphertext}")
    print("  Crib: 'ANGRIFF' at offset 0")
    print("  Searching 17,576 rotor positions...")

    found = bombe_attack(
        ciphertext,
        "ANGRIFF",
        0,
        ("I", "II", "III"),
        verbose=True,
    )
    if found:
        print(f"  ✓ Recovered: {tuple(ALPHA[p] for p in found)}\n")
    else:
        print("  ✗ Not found\n")


def demo_plugboard():
    print("=" * 60)
    print("Plugboard (Steckerbrett) Effect")
    print("=" * 60)
    msg = "SECRETMESSAGE"
    m_no_plug = make_enigma(positions=(0, 0, 0))
    m_with_plug = make_enigma(positions=(0, 0, 0), plugboard_pairs=["AR", "GK", "MT"])

    enc_no = m_no_plug.encrypt(msg)
    enc_yes = m_with_plug.encrypt(msg)

    print(f"  Plaintext:          {msg}")
    print(f"  Without plugboard:  {enc_no}")
    print(f"  With plugboard:     {enc_yes}")
    print("  (Different ciphertext — plugboard adds ~10^14 additional key space)\n")


if __name__ == "__main__":
    demo_enigma_symmetry()
    demo_no_self_encrypt()
    demo_crib_drag()
    demo_bombe_attack()
    demo_plugboard()