[CTF] 布莱切利庄园

import sys
import time
import itertools
import argparse
from multiprocessing import Process, Value, Queue, Event, cpu_count
 
 
def to_int(c): return ord(c) - ord('A')
 
 
def to_char(i): return chr(i + ord('A'))
 
 
def mod26(a): return (a % 26 + 26) % 26
 
 
ROTORS_MAP = {
    0: "EKMFLGDQVZNTOWYHXUSPAIBRCJ",
    1: "AJDKSIRUXBLHWTMCQGZNPYFVOE",
    2: "BDFHJLCPRTXVZNYEIWGAKMUSQO"
}
NOTCHES_MAP = {0: 'Q', 1: 'E', 2: 'V'}
REFLECTOR_B = "YRUHQSLDPXNGOKMIEBFZCWVJAT"
FIXED_ROTOR_INDICES = [0, 1, 2]
FIXED_RINGS = [0, 0, 0]
 
 
class EnigmaMachine:
    def __init__(self, rotor_indices, start_positions, ring_settings, plugboard_pairs=None):
        self.rotor_indices = rotor_indices
        self.initial_positions = list(start_positions)
        self.rings = list(ring_settings)
        self.plugboard = list(range(26))
        if plugboard_pairs:
            pairs_list = plugboard_pairs.upper().strip().split(" ")
            for pair in pairs_list:
                if len(pair) == 2:
                    a = to_int(pair[0]);
                    b = to_int(pair[1])
                    self.plugboard[a] = b
                    self.plugboard[b] = a
 
    def map_forward(self, rotor_map_idx, charIdx, current_positions):
        rotorIdx = self.rotor_indices[rotor_map_idx]
        P = current_positions[rotor_map_idx];
        R = self.rings[rotor_map_idx]
        inputIndex = mod26(charIdx + P - R)
        mappedIndex = to_int(ROTORS_MAP[rotorIdx][inputIndex])
        return mod26(mappedIndex - P + R)
 
    def map_backward(self, rotor_map_idx, charIdx, current_positions):
        rotorIdx = self.rotor_indices[rotor_map_idx]
        P = current_positions[rotor_map_idx];
        R = self.rings[rotor_map_idx]
        inputIndex = mod26(charIdx + P - R)
        charToFind = to_char(inputIndex)
        mappedIndex = ROTORS_MAP[rotorIdx].index(charToFind)
        return mod26(mappedIndex - P + R)
 
    def step(self, current_positions):
        # 这里的逻辑：先判定Notch，再移动
        new_positions = list(current_positions)
        middle_rotor_idx = self.rotor_indices[1]
        right_rotor_idx = self.rotor_indices[2]
 
        isMiddleOnNotch = new_positions[1] == to_int(NOTCHES_MAP[middle_rotor_idx])
        isRightOnNotch = new_positions[2] == to_int(NOTCHES_MAP[right_rotor_idx])
 
        if isMiddleOnNotch:
            new_positions[0] = mod26(new_positions[0] + 1)
            new_positions[1] = mod26(new_positions[1] + 1)
        elif isRightOnNotch:
            new_positions[1] = mod26(new_positions[1] + 1)
 
        new_positions[2] = mod26(new_positions[2] + 1)
        return new_positions
 
    def encrypt(self, input_text, use_plugboard=True):
        result = []
        input_clean = "".join(c for c in input_text.upper() if 'A' <= c <= 'Z')
        current_positions = list(self.initial_positions)
        for c in input_clean:
            current_positions = self.step(current_positions)
            signal = to_int(c)
            if use_plugboard: signal = self.plugboard[signal]
            signal = self.map_forward(2, signal, current_positions)
            signal = self.map_forward(1, signal, current_positions)
            signal = self.map_forward(0, signal, current_positions)
            signal = to_int(REFLECTOR_B[signal])
            signal = self.map_backward(0, signal, current_positions)
            signal = self.map_backward(1, signal, current_positions)
            signal = self.map_backward(2, signal, current_positions)
            if use_plugboard: signal = self.plugboard[signal]
            result.append(to_char(signal))
        return "".join(result)
 
 
def generate_all_plugboard_pairs_2():
    idx = list(range(26))
    for connected in itertools.combinations(idx, 4):
        i1, i2, i3, i4 = connected
        pairings = [((i1, i2), (i3, i4)), ((i1, i3), (i2, i4)), ((i1, i4), (i2, i3))]
        for pairing in pairings:
            yield " ".join([f"{to_char(p[0])}{to_char(p[1])}" for p in pairing])
 
 
def generate_all_plugboard_pairs_1():
    idx = list(range(26))
    for connected in itertools.combinations(idx, 2):
        yield f"{to_char(connected[0])}{to_char(connected[1])}"
 
 
def idx_to_positions(idx):
    a = idx // (26 * 26)
    rem = idx % (26 * 26)
    b = rem // 26
    c = rem % 26
    return (a, b, c)
 
 
def worker_process(start_idx, end_idx, target_ciphertext, plaintext, progress_value: Value, done_event: Event,
                   out_queue: Queue, prefix_len=6):
    plug1 = list(generate_all_plugboard_pairs_1())
    plug2 = list(generate_all_plugboard_pairs_2())
 
    target_prefix = target_ciphertext[:prefix_len]
    plain_prefix = plaintext[:prefix_len]
 
    try:
        for idx in range(start_idx, end_idx):
            if done_event.is_set():
                break
 
            pos = idx_to_positions(idx)
            rotor_positions = [pos[0], pos[1], pos[2]]
 
            # --- 0 对 ---
            m = EnigmaMachine(FIXED_ROTOR_INDICES, rotor_positions, FIXED_RINGS, "")
            if m.encrypt(plain_prefix) == target_prefix:
                if m.encrypt(plaintext) == target_ciphertext:
                    next_pos = m.step(rotor_positions)
                    out_queue.put({
                        'raw_pos': rotor_positions,
                        'next_pos': next_pos,
                        'plugs': "无",
                    })
                    done_event.set()
                    break
 
            # --- 1 对 ---
            found = False
            for pb in plug1:
                m = EnigmaMachine(FIXED_ROTOR_INDICES, rotor_positions, FIXED_RINGS, pb)
                if m.encrypt(plain_prefix) == target_prefix:
                    if m.encrypt(plaintext) == target_ciphertext:
                        next_pos = m.step(rotor_positions)
                        out_queue.put({
                            'raw_pos': rotor_positions,
                            'next_pos': next_pos,
                            'plugs': pb,
                        })
                        done_event.set()
                        found = True
                        break
            if found: break
 
            # --- 2 对 ---
            for pb in plug2:
                m = EnigmaMachine(FIXED_ROTOR_INDICES, rotor_positions, FIXED_RINGS, pb)
                if m.encrypt(plain_prefix) != target_prefix:
                    continue
                if m.encrypt(plaintext) == target_ciphertext:
                    next_pos = m.step(rotor_positions)
                    out_queue.put({
                        'raw_pos': rotor_positions,
                        'next_pos': next_pos,
                        'plugs': pb,
                    })
                    done_event.set()
                    found = True
                    break
            if found: break
 
            with progress_value.get_lock():
                progress_value.value += 1
 
    except Exception as e:
        out_queue.put({'error': repr(e)})
        done_event.set()
 
 
def proc_runner(chunks, target, plain, prog, done_evt, out_q, prefix_len):
    for s, e in chunks:
        if done_evt.is_set(): break
        worker_process(s, e, target, plain, prog, done_evt, out_q, prefix_len)
 
 
def run_mp(target_ciphertext, plaintext, workers=None, chunk=128, prefix_len=6):
    if workers is None:
        workers = max(1, (cpu_count() or 1) - 1)
 
    total_positions = 26 ** 3
 
    print("=" * 60)
    print(f"开始暴力破解 (找到即停)...")
    print(f"目标密文: {target_ciphertext}")
    print(f"对应明文: {plaintext}")
    print(f"Workers: {workers}, Prefix Check: {prefix_len}")
    print("=" * 60)
 
    progress = Value('I', 0)
    done_event = Event()
    out_q = Queue()
 
    indices = [(i, min(i + chunk, total_positions)) for i in range(0, total_positions, chunk)]
    chunks_per_proc = [[] for _ in range(workers)]
    for i, chunk_range in enumerate(indices):
        chunks_per_proc[i % workers].append(chunk_range)
 
    procs = []
    for i in range(workers):
        p = Process(target=proc_runner,
                    args=(chunks_per_proc[i], target_ciphertext, plaintext, progress, done_event, out_q, prefix_len))
        p.start()
        procs.append(p)
 
    start_time = time.time()
    last_print = 0
 
    try:
        while any(p.is_alive() for p in procs):
            # 检查结果
            while not out_q.empty():
                item = out_q.get()
                if 'error' in item:
                    print(f"\n[ERROR] {item['error']}")
                    done_event.set()
                else:
                    raw = item['raw_pos']
                    nxt = item['next_pos']
 
                    # 显示：把 0-based 转成 1-based 显示
                    start_nums = f"{raw[0] + 1},{raw[1] + 1},{raw[2] + 1}"
                    start_chars = f"{to_char(raw[0])}{to_char(raw[1])}{to_char(raw[2])}"
 
                    print("\n" + "#" * 60)
                    print(f"【!!! 找到正确解 !!!】")
                    print(f"初始位置 (穷举): {start_nums:<10} [{start_chars}]")
                    print(f"插线板  (Plugs): {item['plugs']}")
                    print("#" * 60 + "\n")
                    done_event.set()
 
            # 详细进度条
            now = time.time()
            if now - last_print >= 0.5:
                with progress.get_lock():
                    v = progress.value
                pct = v / total_positions * 100
                elapsed = now - start_time
                speed = v / elapsed if elapsed > 0 else 0.0
 
                sys.stdout.write(
                    f"\r进度: {v}/{total_positions} ({pct:.2f}%) | 耗时: {elapsed:.1f}s | 速度: {speed:.1f} pos/s")
                sys.stdout.flush()
                last_print = now
 
            if done_event.is_set():
                time.sleep(0.1)
                break
            else:
                time.sleep(0.05)
 
    finally:
        for p in procs:
            if p.is_alive():
                p.terminate()
        for p in procs:
            p.join()
 
    elapsed = time.time() - start_time
    with progress.get_lock():
        final_v = progress.value
    print(f"\n\n完成: {final_v}/{total_positions} | 总耗时: {elapsed:.1f}s")
    print("-" * 60)
 
 
if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument('ciphertext')
    p.add_argument('plaintext')
    p.add_argument('--workers', type=int, default=None)
    p.add_argument('--chunk', type=int, default=128)
    p.add_argument('--prefix', type=int, default=6)
 
    args = p.parse_args()
 
    target = args.ciphertext.upper().replace(" ", "")
    plain = args.plaintext.upper().replace(" ", "")
 
    run_mp(target, plain, workers=args.workers, chunk=args.chunk, prefix_len=args.prefix)