#!/usr/bin/env python3

import argparse
import base64
import os
import sys
import hashlib
from ecdsa import SigningKey, VerifyingKey, SECP112r1
from Crypto.Cipher import AES
from Crypto.Util import Counter
from argon2.low_level import hash_secret_raw, Type

def generate_keys(private_key_file):
    sk = SigningKey.generate(curve=SECP112r1)
    vk = sk.get_verifying_key()
    priv_bytes = sk.to_string()
    pub_bytes = b'\x04' + vk.to_string()
    pub_b32 = base64.b32encode(pub_bytes).decode('ascii').lower().rstrip('=')
    if private_key_file == False:
        private_key_file = pub_b32[:16] + ".key"
    with open(private_key_file, 'wb') as f:
        f.write(priv_bytes)
    print("Public key: ", end="")
    print(pub_b32)
    print(f"Private key saved to {private_key_file}")

def load_public_key(b32_pub):
    pad = '=' * ((8 - len(b32_pub) % 8) % 8)
    pub_bytes = base64.b32decode(b32_pub.upper() + pad)
    if pub_bytes[0] != 0x04:
        raise ValueError("Invalid public key format")
    return VerifyingKey.from_string(pub_bytes[1:], curve=SECP112r1)

def load_private_key_file(private_key_file):
    with open(private_key_file, 'rb') as f:
        priv_bytes = f.read()
    return SigningKey.from_string(priv_bytes, curve=SECP112r1)

def derive_shared_key(priv_scalar, other_pub_point, key_bytes=16):
    shared_point = other_pub_point * priv_scalar
    shared_x = shared_point.x()
    shared_x_bytes = shared_x.to_bytes(14, byteorder='big')
    h = hashlib.sha256(shared_x_bytes).digest()
    return h[:key_bytes]

def argon2id_kdf(key_material, salt, hash_len=16):
    return hash_secret_raw(
        secret=key_material,
        salt=salt,
        time_cost=2,
        memory_cost=1048576,
        parallelism=1,
        hash_len=hash_len,
        type=Type.ID
    )

def encrypt_file(recipient_b32, infile, outfile):
    recipient_vk = load_public_key(recipient_b32)
    eph_sk = SigningKey.generate(curve=SECP112r1)
    eph_vk = eph_sk.get_verifying_key()
    shared_key = derive_shared_key(eph_sk.privkey.secret_multiplier, recipient_vk.pubkey.point)
    salt = os.urandom(16)
    final_key = argon2id_kdf(shared_key, salt, hash_len=16)
    if infile:
        with open(infile, 'rb') as f:
            plaintext = f.read()
    else:
        plaintext = sys.stdin.buffer.read()
    nonce = os.urandom(8)
    ctr = Counter.new(64, prefix=nonce, initial_value=0)
    cipher = AES.new(final_key, AES.MODE_CTR, counter=ctr)
    ciphertext = cipher.encrypt(plaintext)
    eph_pub_bytes = b'\x04' + eph_vk.to_string()
    outdata = eph_pub_bytes + salt + nonce + ciphertext
    if outfile:
        with open(outfile, 'wb') as f:
            f.write(outdata)
        print(f"Encryption complete. Output written to {outfile}")
    else:
        sys.stdout.buffer.write(outdata)

def decrypt_file(private_key_file, infile, outfile):
    sk = load_private_key_file(private_key_file)
    if infile:
        data = open(infile, 'rb').read()
    else:
        data = sys.stdin.buffer.read()
    if len(data) < 53:
        raise ValueError("Input file too short or corrupted")
    eph_pub_bytes = data[:29]
    if eph_pub_bytes[0] != 0x04:
        raise ValueError("Invalid ephemeral public key format")
    eph_pub = VerifyingKey.from_string(eph_pub_bytes[1:], curve=SECP112r1)
    salt = data[29:29+16]
    nonce = data[29+16:29+16+8]
    ciphertext = data[29+16+8:]
    shared_key = derive_shared_key(sk.privkey.secret_multiplier, eph_pub.pubkey.point)
    final_key = argon2id_kdf(shared_key, salt, hash_len=16)
    ctr = Counter.new(64, prefix=nonce, initial_value=0)
    cipher = AES.new(final_key, AES.MODE_CTR, counter=ctr)
    plaintext = cipher.decrypt(ciphertext)
    if outfile:
        with open(outfile, 'wb') as f:
            f.write(plaintext)
        print(f"Decryption complete. Output written to {outfile}")
    else:
        sys.stdout.buffer.write(plaintext)

def generate_public(private_key_file):
    sk = load_private_key_file(private_key_file)
    vk = sk.get_verifying_key()
    pub_bytes = b'\x04' + vk.to_string()
    pub_b32 = base64.b32encode(pub_bytes).decode('ascii').lower().rstrip('=')
    print("Public key: ", end="")
    print(pub_b32)

def main():
    parser = argparse.ArgumentParser(description="ECC Key Generation and ECIES-like Encryption/Decryption using secp112r1 with Argon2id KDF")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-g", "--generate", action="store_true")
    group.add_argument("-r", "--recipient", type=str)
    group.add_argument("-k", "--private", type=str)
    group.add_argument("-p", "--public", type=str)
    parser.add_argument("-i", "--infile", type=str)
    parser.add_argument("-o", "--outfile", type=str)
    parser.add_argument("-f", "--force", action="store_true")
    args = parser.parse_args()
    if args.generate:
        if not args.outfile:
            args.outfile = False
        if not args.force:
            if args.outfile != False:
                if os.path.exists(args.outfile):
                    print("Output file exists. Use -f to overwrite", file=sys.stderr)
                    sys.exit(1)
        generate_keys(args.outfile)
    elif args.recipient:
        if not args.infile and sys.stdin.isatty():
            print("Encryption mode requires input", file=sys.stderr)
            sys.exit(1)
        infile = args.infile if args.infile else None
        encrypt_file(args.recipient, infile, args.outfile)
    elif args.private:
        if not args.infile and sys.stdin.isatty():
            print("Decryption mode requires input", file=sys.stderr)
            sys.exit(1)
        infile = args.infile if args.infile else None
        decrypt_file(args.private, infile, args.outfile)
    elif args.public:
        generate_public(args.public)

if __name__ == "__main__":
    main()