from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from authorizedkeys.parser import parse_file, dump_file, AuthorizedKey
import sys
from base64 import b64decode, b64encode

# TODO: argument parsing!

# TODO: We currently do not care for authenticity, since we are only concerned
# with the comment. We could sign the rest of the data in the comment and do
# some AEAD, but we currently do not.
key = open('secret', 'rb').read(16)
iv = b"WTF I don't care"
cipher = AES.new(key, AES.MODE_CBC, iv=iv)

output = sys.stdout
decrypt = True if sys.argv[1] == 'decrypt' else False
encrypt = not decrypt
input = open(sys.argv[2]) if len(sys.argv) >= 3 else sys.stdin

# FIXME: file closing

keys = parse_file(input)
for k in keys:
	if isinstance(k, AuthorizedKey):
		if encrypt:
			k.comment = b64encode(cipher.encrypt(pad(k.comment.encode(), cipher.block_size))).decode()
		else: # And now this is just wow.
			k.comment = unpad(cipher.decrypt(b64decode(k.comment)), cipher.block_size).decode()
dump_file(keys, output)