"""
EXAMPLE: ASIS Finals 2018 — "Made by baby" (Theory, 141p / 29 solves)
https://github.com/p4-team/ctf/tree/master/2018-11-24-asis-finals/crypto_baby

VULN: hidden-base knapsack. The plaintext `num` (a PNG) has binary bits b_i
(LSB-first). The ciphertext is
    C = sum_i b_i*(exp^i + (-1)^i) = (sum_i b_i*exp^i) + S
where S = sum_i b_i*(-1)^i is a tiny correction. So (C - S) written in base
`exp` has digits exactly equal to b_i (each 0/1). Once we recover b_i we rebuild
the PNG as  num = sum_i b_i * 2^i  (binary). Only the low bytes of the PNG are
disturbed by the unknown <=512-bit `key`; the header survives.

Run:
    cd /home/code/ctfkit
    python3 examples/crypto_baby/solve.py
"""
from Crypto.Util.number import bytes_to_long, long_to_bytes

HERE = __file__.rsplit("/", 1)[0]


def main():
    data = open(f"{HERE}/flag.enc", "rb").read()
    C = bytes_to_long(data)
    print(f"C bits = {C.bit_length()}")

    for base in range(2, 40):
        for S in range(-200, 201):
            V = C - S
            if V <= 0:
                continue
            # digits of V in `base`, LSB-first, all must be 0/1
            bits = []
            v = V
            ok = True
            n = 0
            while v > 0:
                d = v % base
                if d not in (0, 1):
                    ok = False
                    break
                bits.append(d)
                v //= base
                n += 1
                if n > 40000:
                    ok = False
                    break
            if not ok or not bits:
                continue
            # consistency: rebuild from bits
            if sum(b * (base ** i) for i, b in enumerate(bits)) != V:
                continue
            # rebuild PNG as binary of the same bits
            num = 0
            for i, b in enumerate(bits):
                num += b * (2 ** i)
            out = long_to_bytes(num)
            if out[:8] == b"\x89PNG\r\n\x1a\n":
                print(f"FOUND base={base} S={S}")
                open(f"{HERE}/flag_out.png", "wb").write(out)
                printable = bytes(c if 32 <= c < 127 else 0x20 for c in out)
                print("decoded head:", printable[:300])
                return out
    print("no base found in 2..39 / S in [-200,200]")
    return None


if __name__ == "__main__":
    main()