import os
import secrets
import sys

import compact_xchapoly
import compact_ecdh_curve25519

def keygen():
	sk = os.urandom(32)
	pk = compact_ecdh_curve25519.pubkey(sk)
	return sk, pk

def enc_pk(pk, plaintext):
	ephemeral_sk, ephemeral_pk = keygen()
	# Based on monocypher's
	raw_shared_secret = compact_ecdh_curve25519.ecdh(pk, ephemeral_sk)
	shared_secret = compact_xchapoly.hchacha20(raw_shared_secret, b'\x00'*24)
	nonce = os.urandom(24)
	return ephemeral_pk + nonce + compact_xchapoly.enc(b'', shared_secret, nonce, plaintext)

def dec_sk(sk, ciphertext):
	ephemeral_pk = ciphertext[:32]
	nonce = ciphertext[32:56]
	ciphertext = ciphertext[56:]
	raw_shared_secret = compact_ecdh_curve25519.ecdh(ephemeral_pk, sk)
	shared_secret = compact_xchapoly.hchacha20(raw_shared_secret, b'\x00'*24)
	return compact_xchapoly.dec(b'', shared_secret, nonce, ciphertext)

def usage():
	name=os.path.basename(sys.argv[0])
	print('Usage: %s -G seckey\n       %s -E pubkey\n       %s -D seckey' % (name, name, name), file=sys.stderr)
	sys.exit(1)

if __name__ == '__main__':
	if len(sys.argv) != 3:
		usage()

	if sys.argv[1] == '-G':
		sk, pk = keygen()
		with open(sys.argv[2], 'wb') as f:
			f.write(sk)
		print('pubkey:', pk.hex())
	elif sys.argv[1] == '-E':
		pk = bytes.fromhex(sys.argv[2])
		plaintext = sys.stdin.buffer.read()
		sys.stdout.buffer.write(enc_pk(pk, plaintext))
	elif sys.argv[1] == '-D':
		with open(sys.argv[2], 'rb') as f:
			sk = f.read()
		ciphertext = sys.stdin.buffer.read()
		plaintext = dec_sk(sk, ciphertext)
		if plaintext is None:
			print('%s: Error: Ciphertext authentication failure' % os.path.basename(sys.argv[0]), file=sys.stderr)
			sys.exit(1)
		else:
			sys.stdout.buffer.write(plaintext)
	else:
		usage()