#!/usr/bin/env python3
import argparse
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from elftools.elf.elffile import ELFFile


A = 0x0DC5
B = 0x3DE2


def rot13_byte(x: int) -> int:
    if 0x41 <= x <= 0x5A:
        return ((x - 0x41 + 13) % 26) + 0x41
    if 0x61 <= x <= 0x7A:
        return ((x - 0x61 + 13) % 26) + 0x61
    return x


def mask_for_byte(x: int) -> int:
    signed = x if x < 0x80 else x - 0x100
    return (signed * A + B) & 0xFFFF


def load_text_section(binary_path: str) -> Tuple[int, bytes]:
    with open(binary_path, "rb") as f:
        elf = ELFFile(f)
        text = elf.get_section_by_name(".text")
        if text is None:
            raise RuntimeError("ELF has no .text section")
        return int(text["sh_addr"]), text.data()


def u16_le(buf: bytes, off: int) -> int:
    return buf[off] | (buf[off + 1] << 8)


def call_rel32_bytes(callsite_addr: int, target_addr: int) -> bytes:
    rel = (target_addr - (callsite_addr + 5)) & 0xFFFFFFFF
    return bytes([0xE8]) + rel.to_bytes(4, "little")


@dataclass(frozen=True)
class Constraint:
    block_index: int
    word_index: int
    value: int


def recover_flag(binary_path: str) -> str:
    # These are stable in this challenge binary.
    encrypted_start = 0x1500
    fail_fn = 0x1200

    text_addr, text = load_text_section(binary_path)
    if encrypted_start < text_addr or encrypted_start >= text_addr + len(text):
        raise RuntimeError("encrypted region not inside .text")

    cipher = text[encrypted_start - text_addr :]

    # Precompute reverse maps mask -> byte for the two parities.
    # Parity detail: the SIGILL handler applies ROT13 to the input buffer
    # each time it's invoked, so key bytes alternate between ROT13(flag[i])
    # and flag[i] on each 0x40-aligned UD2 block.
    even_mask_to_byte: Dict[int, int] = {}
    odd_mask_to_byte: Dict[int, int] = {}
    for b in range(256):
        even_mask_to_byte[mask_for_byte(b)] = b
        odd_mask_to_byte[mask_for_byte(rot13_byte(b))] = b

    constraints: Dict[int, int] = {}

    # Scan for places where the decrypted bytes could be:
    #   E8 <rel32-to-0x1200>
    # and where those bytes fully cover at least one 16-bit word.
    for p in range(0, len(cipher) - 5):
        callsite = encrypted_start + p
        seq = call_rel32_bytes(callsite, fail_fn)

        # Map (block,word) -> {0: low_byte, 1: high_byte}
        by_word: Dict[Tuple[int, int], Dict[int, int]] = {}
        for j, plain_b in enumerate(seq):
            o = p + j
            block = o // 0x40
            within = o % 0x40
            word_index = within // 2
            lohi = within % 2
            by_word.setdefault((block, word_index), {})[lohi] = plain_b

        solved: List[Constraint] = []
        ok = True
        for (block, word_index), parts in by_word.items():
            if 0 not in parts or 1 not in parts:
                continue
            plain_w = parts[0] | (parts[1] << 8)
            cipher_w = u16_le(cipher, block * 0x40 + word_index * 2)
            need_mask = cipher_w ^ plain_w

            # Block 0 is the first SIGILL stage => odd parity (ROT13 applied).
            parity_is_odd = (block % 2) == 0
            key_byte = (odd_mask_to_byte if parity_is_odd else even_mask_to_byte).get(need_mask)
            if key_byte is None:
                ok = False
                break
            solved.append(Constraint(block, word_index, key_byte))

        if not ok or not solved:
            continue

        for c in solved:
            prev = constraints.get(c.word_index)
            if prev is not None and prev != c.value:
                raise RuntimeError(
                    f"conflicting constraints for flag[{c.word_index}]: {prev:#x} vs {c.value:#x}"
                )
            constraints[c.word_index] = c.value

    # In this binary, the CALL-to-fail cribs recover indices 5..31.
    # Indices 0..4 are the fixed prefix "TSGCT".
    flag_bytes: List[Optional[int]] = [None] * 32
    for i, ch in enumerate(b"TSGCT"):
        flag_bytes[i] = ch

    for idx, val in constraints.items():
        if 0 <= idx < 32:
            flag_bytes[idx] = val

    if any(v is None for v in flag_bytes):
        missing = [i for i, v in enumerate(flag_bytes) if v is None]
        raise RuntimeError(f"could not recover all bytes; missing indices: {missing}")

    return bytes(flag_bytes).decode("latin1")


def main() -> None:
    ap = argparse.ArgumentParser(description="Recover the medicine flag offline")
    ap.add_argument(
        "--bin",
        default="medicine/medicine",
        help="Path to the medicine ELF (default: medicine/medicine)",
    )
    args = ap.parse_args()

    flag = recover_flag(args.bin)
    print(flag)


if __name__ == "__main__":
    main()