#!/bin/env python

import random
from math import gcd

# the modulus for the lcg prng
n = 256
# the header of the pdf v1.5 file
header = [0x25, 0x50, 0x44, 0x46, 0x2D, 0x31, 0x2E]


def get_factors(ct: bytes, n: int = 256) -> (int, int):
    """find a and b for n and ct"""
    # first generate a list of all numbers without common divisor with 256
    nogcds = [x for x in range(1, n) if gcd(x, n) == 1]
    for b in range(1, n):
        for a in nogcds:
            for i, enc in enumerate(header):
                candidate = (a * enc + b) % n
                if candidate != ct[i]:
                    break
                if i + 1 == len(header):
                    return (a, b)
    print("[-] No solution found.")
    return (0, 0)


def generate_lookuptable(a: int, b: int) -> dict[int, int]:
    """generate a lookuptable for the translation of the enc file"""
    out = {}
    for i in range(0, 256):
        lt = (i * a + b) % 256
        out[lt] = i
    return out


def decrypt(ct: bytes, lut: dict) -> bytes:
    """decrypt the file using the lookup table"""
    res = b""
    for byte in ct:
        dec = lut[byte]
        res += bytes([dec])
    return res


if __name__ == "__main__":
    ct = open("encrypted.bin", "rb").read()
    a, b = get_factors(ct)
    print(f"[+] Factors solved: a={a}, b={b}")
    lut = generate_lookuptable(a, b)
    print(f"[+] Lookuptable generated: \nlookup = {lut}")
    dec = decrypt(ct, lut)
    print(f"[+] Encrypted file decrypted")
    with open("letter.pdf", "wb") as f:
        f.write(dec)