#!/usr/bin/env python3
"""
Firmware encryption tool for Donghui ORPC OTA system.

Usage:
    python3 encrypt_tool.py encrypt --input firmware.bin --output firmware.enc
    python3 encrypt_tool.py decrypt --input firmware.enc --output firmware.bin
    python3 encrypt_tool.py keygen --output keys/
    python3 encrypt_tool.py pack --input firmware.bin --key keys/ --output firmware.ota
"""

import os
import sys
import argparse
import json
import struct
from pathlib import Path

try:
    from Crypto.Cipher import AES
    from Crypto.Hash import SHA256, HMAC
    from Crypto.PublicKey import ECC
    from Crypto.Signature import DSS
except ImportError:
    print("ERROR: pycryptodome not installed. Run: pip install pycryptodome")
    sys.exit(1)


# ==================== Constants ====================

AES_KEY_SIZE = 16       # AES-128
HMAC_KEY_SIZE = 32      # SHA-256 HMAC
IV_SIZE = 16            # AES-CBC IV
HMAC_TAG_SIZE = 32      # SHA-256 HMAC tag
SIGNATURE_SIZE = 64     # ECDSA P-256 signature (DER, but raw = 64 bytes)
FILE_MAGIC = b'ORPC'    # 4-byte magic
FILE_VERSION = 0x0001   # Format version


# ==================== Key Management ====================

def generate_keys(output_dir):
    """Generate AES key, HMAC key, and ECDSA P-256 key pair."""
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    # AES-128 key
    aes_key = os.urandom(AES_KEY_SIZE)
    with open(out / 'aes_key.bin', 'wb') as f:
        f.write(aes_key)
    print(f"  AES-128 key -> {out / 'aes_key.bin'}")

    # HMAC-SHA256 key
    hmac_key = os.urandom(HMAC_KEY_SIZE)
    with open(out / 'hmac_key.bin', 'wb') as f:
        f.write(hmac_key)
    print(f"  HMAC key     -> {out / 'hmac_key.bin'}")

    # ECDSA P-256 key pair
    private_key = ECC.generate(curve='P-256')
    with open(out / 'ecdsa_private.pem', 'wt') as f:
        f.write(private_key.export_key(format='PEM'))
    print(f"  ECDSA priv   -> {out / 'ecdsa_private.pem'}")

    public_key = private_key.public_key()
    with open(out / 'ecdsa_public.pem', 'wt') as f:
        f.write(public_key.export_key(format='PEM'))
    print(f"  ECDSA pub    -> {out / 'ecdsa_public.pem'}")

    # C header for MCU (public key in raw format for OTP)
    pub_bytes = public_key.export_key(format='raw')
    c_header = f"""/*
 * Auto-generated keys for MCU OTP storage.
 * DO NOT EDIT. Regenerate with: python3 encrypt_tool.py keygen
 */

#ifndef _FIRMWARE_KEYS_H_
#define _FIRMWARE_KEYS_H_

#include <stdint.h>

/* AES-128 key for firmware decryption */
static const uint8_t AES_KEY[16] = {{ {', '.join(f'0x{b:02X}' for b in aes_key)} }};

/* HMAC-SHA256 key for integrity verification */
static const uint8_t HMAC_KEY[32] = {{ {', '.join(f'0x{b:02X}' for b in hmac_key)} }};

/* ECDSA P-256 public key (raw 64 bytes: X || Y) */
static const uint8_t ECDSA_PUB_KEY[64] = {{ {', '.join(f'0x{b:02X}' for b in pub_bytes)} }};

#endif /* _FIRMWARE_KEYS_H_ */
"""
    with open(out / 'firmware_keys.h', 'w') as f:
        f.write(c_header)
    print(f"  MCU keys hdr -> {out / 'firmware_keys.h'}")
    print("\nDONE: Keys generated. Keep ecdsa_private.pem SECRET!")
    print("      Store firmware_keys.h in the bootloader project inc/ directory.")
    print("      Upload ecdsa_public.pem to OTA server for signature verification.")


# ==================== Crypto Helpers ====================

def load_keys(key_dir):
    """Load AES key, HMAC key, and private key from directory."""
    key_path = Path(key_dir)

    try:
        with open(key_path / 'aes_key.bin', 'rb') as f:
            aes_key = f.read(AES_KEY_SIZE)
    except FileNotFoundError:
        aes_key = None

    try:
        with open(key_path / 'hmac_key.bin', 'rb') as f:
            hmac_key = f.read(HMAC_KEY_SIZE)
    except FileNotFoundError:
        hmac_key = None

    try:
        with open(key_path / 'ecdsa_private.pem', 'rt') as f:
            private_key = ECC.import_key(f.read())
    except FileNotFoundError:
        private_key = None

    try:
        with open(key_path / 'ecdsa_public.pem', 'rt') as f:
            public_key = ECC.import_key(f.read())
    except FileNotFoundError:
        public_key = None

    return aes_key, hmac_key, private_key, public_key


def aes_cbc_encrypt(data, key, iv=None):
    """AES-128-CBC encrypt with PKCS7 padding."""
    if iv is None:
        iv = os.urandom(IV_SIZE)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    # PKCS7 padding
    pad_len = AES.block_size - (len(data) % AES.block_size)
    padded = data + bytes([pad_len] * pad_len)
    encrypted = cipher.encrypt(padded)
    return iv + encrypted


def aes_cbc_decrypt(data, key):
    """AES-128-CBC decrypt with PKCS7 unpadding."""
    iv = data[:IV_SIZE]
    encrypted = data[IV_SIZE:]
    cipher = AES.new(key, AES.MODE_CBC, iv)
    padded = cipher.decrypt(encrypted)
    # PKCS7 unpad
    pad_len = padded[-1]
    if pad_len < 1 or pad_len > AES.block_size:
        raise ValueError("Invalid PKCS7 padding")
    return padded[:-pad_len]


def hmac_sha256(data, key):
    """Compute HMAC-SHA256."""
    h = HMAC.new(key, digestmod=SHA256)
    h.update(data)
    return h.digest()


def ecdsa_sign(data, private_key):
    """ECDSA P-256 sign using deterministic signature."""
    h = SHA256.new(data)
    signer = DSS.new(private_key, 'fips-186-3')
    return signer.sign(h)


def ecdsa_verify(data, signature, public_key):
    """Verify ECDSA P-256 signature."""
    h = SHA256.new(data)
    verifier = DSS.new(public_key, 'fips-186-3')
    try:
        verifier.verify(h, signature)
        return True
    except (ValueError, TypeError):
        return False


# ==================== File Format ====================

# Encrypted firmware file (.enc):
# +--------+--------+--------+--------------+-------------+
# | IV(16) | AES-Encrypted Firmware | HMAC-Tag(32) |
# +--------+--------+--------+--------------+-------------+

# Full OTA package (.ota):
# +----------+--------+--------+--------+--------+------------+----------+
# | Magic(4) | Ver(2) | Size(4) | CRC32(4) | IV(16) | Encrypted  | HMAC(32) | Signature(64) |
# | "ORPC"   | 0x0001 | payload | of plain  |        | firmware   | SHA256   | ECDSA P-256    |
# +----------+--------+--------+--------+--------+------------+----------+


def create_encrypted_file(firmware_bin, aes_key, hmac_key):
    """Create .enc file (IV + encrypted + HMAC)."""
    # Generate random IV
    iv = os.urandom(IV_SIZE)

    # AES-CBC encrypt
    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
    pad_len = AES.block_size - (len(firmware_bin) % AES.block_size)
    padded = firmware_bin + bytes([pad_len] * pad_len)
    encrypted = cipher.encrypt(padded)

    # HMAC-SHA256 over (IV + encrypted)
    h = HMAC.new(hmac_key, digestmod=SHA256)
    h.update(iv + encrypted)
    tag = h.digest()

    return iv + encrypted + tag


def parse_encrypted_file(enc_data, aes_key, hmac_key):
    """Parse .enc file, verify HMAC, decrypt."""
    if len(enc_data) < IV_SIZE + HMAC_TAG_SIZE:
        raise ValueError("File too small")

    iv = enc_data[:IV_SIZE]
    tag = enc_data[-HMAC_TAG_SIZE:]
    encrypted = enc_data[IV_SIZE:-HMAC_TAG_SIZE]

    # Verify HMAC first
    h = HMAC.new(hmac_key, digestmod=SHA256)
    h.update(iv + encrypted)
    expected_tag = h.digest()
    if tag != expected_tag:
        raise ValueError("HMAC verification FAILED! File may be corrupted or tampered.")

    # Decrypt
    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
    padded = cipher.decrypt(encrypted)
    # Unpad
    pad_len = padded[-1]
    if pad_len < 1 or pad_len > AES.block_size:
        raise ValueError("Invalid PKCS7 padding")
    return padded[:-pad_len]


def create_ota_package(firmware_bin, aes_key, hmac_key, private_key,
                       version=FILE_VERSION):
    """Create full .ota package with signature."""
    # Encrypt
    payload = create_encrypted_file(firmware_bin, aes_key, hmac_key)

    # Compute CRC32 of plaintext
    crc32 = struct.pack('<I', binascii.crc32(firmware_bin) & 0xFFFFFFFF)

    # Compute signature over (magic + ver + size + crc32 + payload)
    header = FILE_MAGIC
    header += struct.pack('<H', version)
    header += struct.pack('<I', len(firmware_bin))
    header += crc32

    data_to_sign = header + payload
    signature = ecdsa_sign(data_to_sign, private_key)

    return data_to_sign + signature


# ==================== CLI ====================

def cmd_keygen(args):
    print("Generating OTA encryption keys...")
    generate_keys(args.output)


def cmd_encrypt(args):
    print(f"Encrypting: {args.input}")
    aes_key, hmac_key, priv_key, _ = load_keys(args.key_dir)
    if aes_key is None or hmac_key is None:
        print("ERROR: Missing AES or HMAC key. Run 'keygen' first.")
        sys.exit(1)

    with open(args.input, 'rb') as f:
        firmware = f.read()

    enc_data = create_encrypted_file(firmware, aes_key, hmac_key)

    with open(args.output, 'wb') as f:
        f.write(enc_data)

    print(f"  Input:  {len(firmware)} bytes")
    print(f"  Output: {len(enc_data)} bytes ({args.output})")
    print("  Encrypted with AES-128-CBC + HMAC-SHA256.")


def cmd_decrypt(args):
    print(f"Decrypting: {args.input}")
    aes_key, hmac_key, _, _ = load_keys(args.key_dir)
    if aes_key is None or hmac_key is None:
        print("ERROR: Missing AES or HMAC key.")
        sys.exit(1)

    with open(args.input, 'rb') as f:
        enc_data = f.read()

    try:
        plain = parse_encrypted_file(enc_data, aes_key, hmac_key)
    except ValueError as e:
        print(f"ERROR: {e}")
        sys.exit(1)

    with open(args.output, 'wb') as f:
        f.write(plain)

    print(f"  Input:  {len(enc_data)} bytes")
    print(f"  Output: {len(plain)} bytes ({args.output})")
    print("  HMAC verified. Decryption OK.")


def cmd_pack(args):
    print(f"Creating OTA package: {args.input}")
    aes_key, hmac_key, priv_key, _ = load_keys(args.key_dir)
    if aes_key is None or hmac_key is None or priv_key is None:
        print("ERROR: Missing keys (need AES, HMAC, and ECDSA private key).")
        sys.exit(1)

    with open(args.input, 'rb') as f:
        firmware = f.read()

    ota_data = create_ota_package(firmware, aes_key, hmac_key, priv_key)

    with open(args.output, 'wb') as f:
        f.write(ota_data)

    print(f"  Firmware: {len(firmware)} bytes")
    print(f"  Package:  {len(ota_data)} bytes ({args.output})")
    print("  AES-128-CBC encrypted + HMAC authenticated + ECDSA signed.")
    print(f"  OTA package ready for distribution.")


# ==================== Main ====================

if __name__ == '__main__':
    import binascii  # for CRC32 in pack command

    parser = argparse.ArgumentParser(
        description='Donghui ORPC OTA Firmware Encryption Tool')
    sub = parser.add_subparsers(dest='command', required=True)

    # keygen
    kg = sub.add_parser('keygen', help='Generate encryption keys')
    kg.add_argument('--output', '-o', default='./keys',
                    help='Output directory for keys')

    # encrypt
    en = sub.add_parser('encrypt', help='Encrypt firmware binary')
    en.add_argument('--input', '-i', required=True, help='Raw firmware .bin')
    en.add_argument('--output', '-o', required=True, help='Output .enc file')
    en.add_argument('--key-dir', '-k', default='./keys',
                    help='Key directory')

    # decrypt
    de = sub.add_parser('decrypt', help='Decrypt firmware for testing')
    de.add_argument('--input', '-i', required=True, help='Encrypted .enc file')
    de.add_argument('--output', '-o', required=True, help='Output .bin file')
    de.add_argument('--key-dir', '-k', default='./keys',
                    help='Key directory')

    # pack (full OTA package with signature)
    pk = sub.add_parser('pack', help='Create full OTA package')
    pk.add_argument('--input', '-i', required=True, help='Raw firmware .bin')
    pk.add_argument('--output', '-o', required=True, help='Output .ota file')
    pk.add_argument('--key-dir', '-k', default='./keys',
                    help='Key directory')

    args = parser.parse_args()

    if args.command == 'keygen':
        cmd_keygen(args)
    elif args.command == 'encrypt':
        cmd_encrypt(args)
    elif args.command == 'decrypt':
        cmd_decrypt(args)
    elif args.command == 'pack':
        cmd_pack(args)