import base64
import struct
import nacl.secret
import nacl.utils
import nacl.hashlib
import nacl.hash

BUFFER_SIZE = 4 * (1024 * 1024)

def read_file_blocks(file, extra_bytes=0):
    while True:
        data = file.read(BUFFER_SIZE + extra_bytes)
        if not data:
            break
        yield data

def hmac_file(file, key):
    blake = nacl.hashlib.blake2b(key=key)
    with open(file, 'rb') as in_file:
        for block in read_file_blocks(in_file):
            blake.update(block)
    return blake.hexdigest()

def encrypt_archive(archive_name, encrypted_name):
    key = nacl.utils.random(nacl.secret.SecretBox.KEY_SIZE)
    #Use 4 bytes less than the nonce size to make room for the block counter
    nonce = nacl.utils.random(nacl.secret.SecretBox.NONCE_SIZE - 4)
    block_num = 0

    box = nacl.secret.SecretBox(key)
    with open(archive_name, 'rb') as in_file, open(encrypted_name, 'wb') as out_file:
        for data in read_file_blocks(in_file):
            #Append the block counter to the nonce, so each block has a unique nonce
            block_nonce = nonce + struct.pack(">I", block_num)
            block = box.encrypt(data, block_nonce)
            out_file.write(block.ciphertext)
            block_num += 1

    hmac_key = nacl.hash.sha256(key + nonce, encoder=nacl.encoding.RawEncoder)
    output = {}
    output['key'] = base64.b64encode(key + nonce)
    output['signature'] = hmac_file(encrypted_name, hmac_key)
    return output

def decrypt_archive(encrypted_name, archive_name, key_info):
    key_bytes = base64.b64decode(key_info['key'])

    key = key_bytes[:nacl.secret.SecretBox.KEY_SIZE]
    nonce = key_bytes[nacl.secret.SecretBox.KEY_SIZE:]

    extra_bytes = nacl.secret.SecretBox.MACBYTES
    hmac_key = nacl.hash.sha256(key_bytes, encoder=nacl.encoding.RawEncoder)
    hmac = hmac_file(encrypted_name, hmac_key)
    if hmac != key_info['signature']:
        print('hmac mismatch')
        return

    block_num = 0
    box = nacl.secret.SecretBox(key)
    with open(encrypted_name, 'rb') as in_file, open(archive_name, 'wb') as out_file:
        # nacl adds a MAC to each block, when reading the file in, this needs to be taken into account
        for data in read_file_blocks(in_file, extra_bytes=extra_bytes):
            block_nonce = nonce + struct.pack(">I", block_num)
            block = box.decrypt(data, block_nonce)
            out_file.write(block)
            block_num += 1

key_info = encrypt_archive("C:\temp\test.csv", "C:\temp\test.enc")
print(key_info)
decrypt_archive("C:\temp\test.enc", "C:\temp\test.enc.csv", key_info)