#!/usr/local/bin/python
import hashlib
from os import urandom
from flag import FLAG


def gen_pubkey(secret: bytes, hasher=hashlib.sha512) -> list:
    def hash(m):
        return hasher(m).digest()

    state = hash(secret)
    pubkey = []
    for _ in range(len(hash(b"0")) * 4):
        pubkey.append(int.from_bytes(state, "big"))
        state = hash(state)
    return pubkey


def happiness(x: int) -> int:
    return x - sum((x >> i) for i in range(1, x.bit_length()))


def encode_message(message: bytes, segment_len: int) -> list:
    message += bytes(segment_len - len(message) % (segment_len))
    padded = []
    for i in range(0, len(message), segment_len):
        block = message[i : i + segment_len]
        padded.append(urandom(segment_len * 3) + block + urandom(segment_len * 4))
    return [int.from_bytes(i, "big") for i in padded]


def encrypt(pubkey: list, message: bytes) -> list:
    encrypted_blocks = []
    for block_int in encode_message(message, len(pubkey) // 32):
        encrypted_blocks.append([happiness(i & block_int) for i in pubkey])
    return encrypted_blocks


secret = urandom(16)
A = gen_pubkey(secret)
enc = encrypt(A, FLAG)
print(secret.hex())
print(enc)