from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import marshal
import os

aes_key = os.urandom(17)
with open('../obfuscator/build/lib.win-amd64-cpython-310/pyhumor_runtime.cp310-win_amd64.pyd', 'rb') as fin:
    filled = fin.read().replace(b'WAIT_TO_BE_FILLED', aes_key)
with open('pyhumor_runtime.cp310-win_amd64.pyd', 'wb') as fout:
    fout.write(filled)

code = compile(open('plain.py').read(), "<frozen>", "exec")
raw = marshal.dumps(code)

code_len = len(code.co_code)
encrypt_len = (code_len - 8) - (code_len - 8) % 16
cipher = AES.new(aes_key[1:], AES.MODE_ECB)
encrypted = cipher.encrypt(raw[38:38+encrypt_len])
assert len(encrypted) == encrypt_len
raw = raw[:38] + encrypted + raw[38+encrypt_len:]
raw = pad(raw, 16)

cipher = AES.new(aes_key[:16], AES.MODE_ECB)
encrypted = cipher.encrypt(raw)
structured = b'PYHUMOR\x00\x00\x03\x0A\x00' + len(encrypted).to_bytes(4, 'little') + encrypted

open("pyhumor.py", "w").write(f'''from pyhumor_runtime import __pyhumor__

__pyhumor__(__name__, __file__, {repr(structured)})
''')