import math
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

p = 6149059152047173810929520610252054724434547268293997900228346988632026139296516495540477511197233410835031806302501788817802147338507673010790914591966411034143018465230753723388103315997074066465915406376864167374655977885785821310570662894392390395753471445736267810187903833304863045974238749688245530607
g = 5719275944239987190458421245546229548888846731353398043162644430064597096976159819921524452554095058636720054596231559684320789722962976527976545669345058297524867469181858053038816720679526830291841309662435186455700706221003964190466817684147617948365280845302761388758671140064249106209075533637357263231
A = 2555297614863845215431835415772614784743798924133900800787463671501662222468554222725877277092258849851944927347384717791540169865769130437660089795213081276681792643860918393721387336966064482776321260842945637479291021202655103477316865068516592795061612417388401255868653332471670086454998195709403176041
c1 = 5512928361799616712522312337586877320563049808521259994710385873918851237761731656487251356100346025012478907551450133321786614899546347085949956566607223672838137967339204274717066193165057197169297068843124894954077280100769193924635067973901636912604952268847061931872724414381559827589687222654708496705
c2 = 210477289627346949522196934377558219236062280271196462092468487627372551930774629985446433116041864076509672563413245745389355377479670040469926801201354051811351515045242197387054899103549818741787539277377997302764841143680449130728621812269605318109456762177695192938113314934087213735582325586658379425
SMOOTH_FACTORS = [2, 17, 47, 64_245_481, 1_741_212_479_437, 5_733_063_549_001]

def baby_step_giant_step(g: int, h: int, p: int, order: int) -> int | None:
    """ Solve g^x ≡ h (mod p) with x in [0, order). Time & space: O(√order). """
    m = math.isqrt(order) + 1
    table = {}
    val = 1
    for j in range(m):
        table[val] = j
        val = val * g % p
    g_inv_m = pow(pow(g, m, p), p - 2, p)
    val = h
    for i in range(m):
        if val in table:
            candidate = i * m + table[val]
            if candidate < order:
                return candidate
        val = val * g_inv_m % p
    return None # No solution found

def pohlig_hellman(g: int, A: int, p: int, smooth_factors: list[int]) -> tuple[list, list]:
    """ Returns (residues, moduli) such that: a ≡ residues[i] (mod moduli[i]) for each small prime factor. """
    residues = []
    moduli = []
    print("
[+] Pohlig-Hellman: solving DLP in small subgroups...")
    for q in smooth_factors:
        exp = (p - 1) // q
        g_q = pow(g, exp, p) # subgroup generator of order q
        A_q = pow(A, exp, p) # A projected into subgroup
        a_q = baby_step_giant_step(g_q, A_q, p, q)
        assert pow(g_q, a_q, p) == A_q, f"BSGS failed for q={q}"
        print(f" a mod {q:<20} = {a_q}")
        residues.append(a_q)
        moduli.append(q)
    return residues, moduli

def extended_gcd(a: int, b: int) -> tuple[int, int, int]:
    """Returns (gcd, x, y) satisfying a·x + b·y = gcd."""
    if b == 0: return a, 1, 0
    g, x, y = extended_gcd(b, a % b)
    return g, y, x - (a // b) * y

def crt(residues: list[int], moduli: list[int]) -> int:
    """Combine residues via Chinese Remainder Theorem."""
    M = 1
    for m in moduli: M *= m
    x = 0
    for r, m in zip(residues, moduli):
        Mi = M // m
        _, inv, _ = extended_gcd(Mi, m)
        x += r * Mi * inv
    return x % M

def elgamal_decrypt(c1: int, c2: int, a: int, p: int) -> bytes:
    """Decrypt an ElGamal ciphertext given private key a."""
    shared = pow(c1, a, p) # c1^a = g^{ka}
    shared_inv = pow(shared, p - 2, p) # modular inverse via Fermat
    m = c2 * shared_inv % p
    return m.to_bytes((m.bit_length() + 7) // 8, "big")

def aes_cbc_decrypt(key: bytes, ciphertext: bytes, iv: bytes = b"�" * 16) -> bytes:
    """Decrypt AES-CBC ciphertext."""
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    return cipher.decryptor().update(ciphertext)

FLAG_ENC_PATH = "flag.enc" # 48-byte AES-encrypted flag

def main():
    print("=" * 64)
    print(" ElGamal Cryptanalysis — Pohlig-Hellman Attack")
    print("=" * 64)

    # ── Phase 1: Factor p-1 ──────────────────
    print("
[+] Known smooth factors of p-1:")
    smooth_product = 1
    for q in SMOOTH_FACTORS:
        print(f" {q}")
        smooth_product *= q
    print(f"
 Smooth product = {smooth_product}")
    print(f" Smooth product bit-length = {smooth_product.bit_length()} bits")

    # ── Phase 2 & 3: Pohlig-Hellman ──────────
    residues, moduli = pohlig_hellman(g, A, p, SMOOTH_FACTORS)

    # ── Phase 4: CRT ─────────────────────────
    a = crt(residues, moduli)
    print(f"
[+] Private key recovered via CRT:")
    print(f" a = {a}")
    print(f" a bit-length = {a.bit_length()} bits (fits in smooth product!)")

    # Verify
    assert pow(g, a, p) == A, "Private key verification FAILED!"
    print(" ✓ Verified: g^a ≡ A (mod p)")

    # ── Phase 5: ElGamal Decrypt ─────────────
    aes_key = elgamal_decrypt(c1, c2, a, p)
    print(f"
[+] ElGamal decrypted AES key:")
    print(f" {aes_key.hex()}")
    assert len(aes_key) == 16, f"Expected 16-byte AES key, got {len(aes_key)}"

    # ── Phase 6: AES Decrypt flag ─────────────
    try:
        with open(FLAG_ENC_PATH, "rb") as f:
            flag_enc = f.read()
    except FileNotFoundError:
        # Hardcode if file not present
        flag_enc = bytes.fromhex(
            "c79fde239c08f4523727eef9ba933e69"
            "03d9e4f469830b154cc4a32d88b492e7"
            "5cb41a544351c821466153409048a581"
        )
    plaintext = aes_cbc_decrypt(aes_key, flag_enc)

    # Strip PKCS#7 padding
    pad_len = plaintext[-1]
    flag = plaintext[:-pad_len].decode("utf-8", errors="replace")

    print(f"
{'=' * 64}")
    print(f"FLAG: {flag}")
    print(f"{'=' * 64}
")
    return flag

if __name__ == "__main__":
    main()