#!/usr/bin/env python3

load("TD02.sage")

import bitarray
try:
    from tqdm.auto import tqdm
except ImportError:
    def tqdm(x, *args, **kwargs):
        return x

"""
Dans TD02.sage j'ai défini la fonction LFSR_step(P, state).
J'ai aussi choisi un polynôme de rétroaction P mystère.

Indice, FunWithFlags.png est un vrai fichier .png valide et P est de degré 64.
"""

F2 = GF(2)
A.<X> = F2[]


def encrypt(path, P = P):
    N = P.degree()
    state_I = [ F2.random_element() for _ in range(N) ]
    state = copy(state_I)
    with open(path, "rb") as flag:
        data = flag.read()

    bIn = bitarray.bitarray()
    bIn.frombytes(data)

    bOut = bitarray.bitarray()
    for b in tqdm(bIn):
        bit, state = LFSR_step(P, state)
        bOut.append(F2(b) + bit)

    outpath = path + ".enc"

    with open(outpath, "wb") as flag_enc:
        flag_enc.write(bOut.tobytes())

    return state_I


def decrypt(path, state_I, P=P):
    with open(path, "rb") as flag_enc:
        data = flag_enc.read()

    bIn = bitarray.bitarray()
    bIn.frombytes(data)


    bOut = bitarray.bitarray()

    state = state_I

    for b in bIn:
        bit, state = LFSR_step(P, state)
        bOut.append(F2(b) + bit)

    outpath = path.removesuffix(".enc") + ".dec.png"

    with open(outpath, "wb") as flag_dec:
        flag_dec.write(bOut.tobytes())