import secrets

import chacha20
import poly1305

def unleu64(n):
	assert n >> 64 == 0
	return bytes((n >> i) & 0xff for i in range(0, 64, 8))

def pad16(x):
	if len(x) % 16 == 0:
		return b''
	else:
		return b'\x00' * (16 - (len(x) % 16))

def poly1305_key_gen(key, nonce):
	return chacha20.chacha20_block(key, 0, nonce)[:32]

def generate_tag(aad, key, nonce, ciphertext):
	key = poly1305_key_gen(key, nonce)
	message = b''.join((
		aad, pad16(aad),
		ciphertext, pad16(ciphertext),
		unleu64(len(aad)),
		unleu64(len(ciphertext))
	))
	return poly1305.poly1305(message, key)

def chapoly_aead_enc(aad, key, nonce, plaintext):
	ciphertext = bytes(chacha20.chacha20(key, 1, nonce, plaintext))
	tag = generate_tag(aad, key, nonce, ciphertext)
	return ciphertext + tag

def chapoly_aead_dec(aad, key, nonce, ciphertext):
	ciphertext, tag = ciphertext[:-16], ciphertext[-16:]
	expected_tag = generate_tag(aad, key, nonce, ciphertext)
	if not secrets.compare_digest(tag, expected_tag):
		return None
	return bytes(chacha20.chacha20(key, 1, nonce, ciphertext))

if __name__ == '__main__':
	plaintext = b"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it."
	aad = bytes.fromhex('50 51 52 53 c0 c1 c2 c3 c4 c5 c6 c7')
	key = bytes.fromhex('80 81 82 83 84 85 86 87 88 89 8a 8b 8c 8d 8e 8f 90 91 92 93 94 95 96 97 98 99 9a 9b 9c 9d 9e 9f')
	iv = b'@ABCDEFG'
	constant = bytes.fromhex('07 00 00 00')
	nonce = constant + iv
	print(chapoly_aead_enc(aad, key, nonce, plaintext) == bytes.fromhex('d3 1a 8d 34 64 8e 60 db 7b 86 af bc 53 ef 7e c2 a4 ad ed 51 29 6e 08 fe a9 e2 b5 a7 36 ee 62 d6 3d be a4 5e 8c a9 67 12 82 fa fb 69 da 92 72 8b 1a 71 de 0a 9e 06 0b 29 05 d6 a5 b6 7e cd 3b 36 92 dd bd 7f 2d 77 8b 8c 98 03 ae e3 28 09 1b 58 fa b3 24 e4 fa d6 75 94 55 85 80 8b 48 31 d7 bc 3f f4 de f0 8e 4b 7a 9d e5 76 d2 65 86 ce c6 4b 61 16 1a:e1:0b:59:4f:09:e2:6a:7e:90:2e:cb:d0:60:06:91'.replace(':','')))