from Crypto.Util.number import getPrime, isPrime

e = 65537
N = 1
while (N - 1) % e == 0:
    N = getPrime(2048)

def x(a, b):
    return bytes(x^y for x,y in zip(a,b))

class C():
    def __init__(self, N, e):
        self.N = N
        self.e = e
        self._state = b"\x00" * 256
        self.s = set()

    def _h(self, block):
        assert len(block) == 256

        if block in self.s:
            raise ValueError("bad")
        self.s.add(block)

        data = int.from_bytes(block, "big")
        if data < 2 or data >= N-1:
            raise ValueError("bad")

        data = pow(data, self.e, self.N)
        if data in self.s:
            raise ValueError("bad")
        self.s.add(data)

        return data.to_bytes(256, "big")

    def u(self, data):
        assert len(data) % 256 == 0

        for block in range(0, len(data), 256):
            block = data[block:block+256]
            self._state = x(self._state, self._h(block))

        return self

    def hd(self):
        return self._state.hex()

    def __repr__(self):
        return f"C({self.N}, {self.e})"

def m():
    h = C(N, e)
    print(h)

    print("Input the data you'd like to process:")
    p = bytes.fromhex(input("> "))
    if len(p) < 256 or len(p) % 256 != 0:
        raise ValueError("Invalid input!")

    z = h.u(p).hd()
    print("hash(input) == ", z)
    if z == "00" * 256:
        with open("/flag","r") as f:
            print(f.read())
    else:
        print("...")

m()
