#!/usr/bin/env python3
import hashlib, hmac
import binascii, sys, struct, os
from hexdump import hexdump 
from Crypto.Cipher import AES
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, padding, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.serialization import PublicFormat   
from cryptography.hazmat.primitives.serialization import Encoding       

def encrypt_file(key, iv, in_filename, out_filename=None, chunksize=64*1024):
    if not out_filename:
        out_filename = in_filename + '.enc'

    encryptor = AES.new(key, AES.MODE_CBC, iv)
    filesize = os.path.getsize(in_filename)

    with open(in_filename, 'rb') as infile:
        with open(out_filename, 'wb') as outfile:
            outfile.write(struct.pack('<Q', filesize))

            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    chunk += b' ' * (16 - len(chunk) % 16)

                outfile.write(encryptor.encrypt(chunk))

def decrypt_file(key, iv, in_filename, out_filename=None, chunksize=24*1024):
    if not out_filename:
        out_filename = os.path.splitext(in_filename)[0]

    with open(in_filename, 'rb') as infile:
        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
        decryptor = AES.new(key, AES.MODE_CBC, iv)

        with open(out_filename, 'wb') as outfile:
            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                outfile.write(decryptor.decrypt(chunk))

            outfile.truncate(origsize)

# define -----------------------------------------------------------------|
curve = ec.SECP256R1()
signature_algorithm = ec.ECDSA(hashes.SHA256())

salt=b'MOTIECTF-MessageKeys'
hkdf = HKDF(
    algorithm=hashes.SHA256(),
    length=0x30,
    salt=salt,
    info=None,
    backend=default_backend()
)

ephemeralKey     = ec.generate_private_key(curve, default_backend())   # ephemeral_private
ephemeralPubKey  = ephemeralKey.public_key()
print("[+] ephemeralPubKey ---")
hexdump(ephemeralPubKey.public_bytes(Encoding.DER, PublicFormat.SubjectPublicKeyInfo))

privKey     = int.from_bytes(bytes.fromhex("5a55034a6c8ce32e efc745faf7e5e2a8 d24cadd2116ab132 8b634f21f6b21706"), "big")  
privKey     = ec.derive_private_key(privKey, curve, default_backend())
pubKey      = privKey.public_key()

shared_key  = ephemeralKey.exchange(ec.ECDH(), pubKey)
derived_key = hkdf.derive(shared_key)
key         = derived_key[:0x20]
iv          = derived_key[0x20:]

# Encrypting file 
print("[+] Encrypting file")
filename    = 'chal.zip' 
encrypt_file(key, iv, filename) 
 
print("[+] Encryption DONE")