encryption_engine.py

"""High-integrity ECC encryption utilities with secure key handling.

This module implements an authenticated ECIES-style workflow using
elliptic-curve Diffie-Hellman key agreement and AES-GCM for payload
confidentiality and integrity. A lightweight CLI is exposed for key
management, encryption, and decryption flows suitable for developer
workstations and reproducible pipelines.
"""
from __future__ import annotations

import argparse
import base64
import json
import os
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Tuple

try:
    from cryptography.exceptions import InvalidTag
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.hashes import SHA256
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
except ModuleNotFoundError as exc:  # pragma: no cover - import guard
    raise ModuleNotFoundError(
        "The 'cryptography' package is required. Install it via 'pip install cryptography'."
    ) from exc

HKDF_INFO = b"QuantumChaosEncryption-v2"
DEFAULT_CURVE = ec.SECP256R1()
NONCE_LENGTH = 12
SALT_LENGTH = 16


def generate_ecc_keys() -> Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]:
    """Generate an ECC keypair on the default curve."""

    private_key = ec.generate_private_key(DEFAULT_CURVE)
    return private_key, private_key.public_key()


def serialize_public_key(public_key: ec.EllipticCurvePublicKey) -> str:
    """Serialize a public key to base64-encoded uncompressed point bytes."""

    encoded = public_key.public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
    )
    return base64.b64encode(encoded).decode("utf-8")


def deserialize_public_key(data: str) -> ec.EllipticCurvePublicKey:
    """Deserialize a base64-encoded uncompressed point into a public key."""

    raw = base64.b64decode(data.encode("utf-8"))
    return ec.EllipticCurvePublicKey.from_encoded_point(DEFAULT_CURVE, raw)


def serialize_private_key(
    private_key: ec.EllipticCurvePrivateKey, password: Optional[bytes] = None
) -> bytes:
    """Serialize a private key to PEM with optional password protection."""

    encryption_algorithm: serialization.KeySerializationEncryption
    if password:
        encryption_algorithm = serialization.BestAvailableEncryption(password)
    else:
        encryption_algorithm = serialization.NoEncryption()

    return private_key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        encryption_algorithm,
    )


def load_private_key(pem_data: bytes, password: Optional[bytes] = None) -> ec.EllipticCurvePrivateKey:
    """Load a private key from PEM data with the provided password."""

    return serialization.load_pem_private_key(pem_data, password)


@dataclass
class EncryptedPayload:
    """Canonical representation of an authenticated ciphertext bundle."""

    ephemeral_public_key: str
    salt: str
    nonce: str
    ciphertext: str
    associated_data: Optional[str] = None

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2)

    @classmethod
    def from_json(cls, data: str) -> "EncryptedPayload":
        payload = json.loads(data)
        return cls(**payload)


def _derive_aead_key(shared_secret: bytes, salt: bytes) -> bytes:
    hkdf = HKDF(
        algorithm=SHA256(),
        length=32,
        salt=salt,
        info=HKDF_INFO,
    )
    return hkdf.derive(shared_secret)


def encrypt_data(
    recipient_public_key: ec.EllipticCurvePublicKey,
    data: bytes | str,
    *,
    associated_data: Optional[bytes | str] = None,
) -> EncryptedPayload:
    """Encrypt data for the recipient using ephemeral ECDH and AES-GCM."""

    if isinstance(data, str):
        plaintext = data.encode("utf-8")
    else:
        plaintext = data

    if associated_data is not None and isinstance(associated_data, str):
        aad_bytes = associated_data.encode("utf-8")
        aad_repr = associated_data
    else:
        aad_bytes = associated_data
        aad_repr = associated_data.decode("utf-8") if associated_data else None

    ephemeral_private = ec.generate_private_key(DEFAULT_CURVE)
    ephemeral_public = ephemeral_private.public_key()
    shared_secret = ephemeral_private.exchange(ec.ECDH(), recipient_public_key)

    salt = os.urandom(SALT_LENGTH)
    nonce = os.urandom(NONCE_LENGTH)
    key = _derive_aead_key(shared_secret, salt)

    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext, aad_bytes)

    # Best-effort zeroization of derived key material.
    if isinstance(key, bytes):
        mutable_key = bytearray(key)
        for i in range(len(mutable_key)):
            mutable_key[i] = 0

    payload = EncryptedPayload(
        ephemeral_public_key=serialize_public_key(ephemeral_public),
        salt=base64.b64encode(salt).decode("utf-8"),
        nonce=base64.b64encode(nonce).decode("utf-8"),
        ciphertext=base64.b64encode(ciphertext).decode("utf-8"),
        associated_data=aad_repr,
    )
    return payload


def decrypt_data(
    recipient_private_key: ec.EllipticCurvePrivateKey,
    payload: EncryptedPayload,
    *,
    associated_data: Optional[bytes | str] = None,
) -> bytes:
    """Decrypt payload using the recipient's private key and AES-GCM."""

    ephemeral_public = deserialize_public_key(payload.ephemeral_public_key)
    shared_secret = recipient_private_key.exchange(ec.ECDH(), ephemeral_public)

    salt = base64.b64decode(payload.salt.encode("utf-8"))
    nonce = base64.b64decode(payload.nonce.encode("utf-8"))
    ciphertext = base64.b64decode(payload.ciphertext.encode("utf-8"))

    if associated_data is None and payload.associated_data is not None:
        aad_bytes: Optional[bytes] = payload.associated_data.encode("utf-8")
    elif isinstance(associated_data, str):
        aad_bytes = associated_data.encode("utf-8")
    else:
        aad_bytes = associated_data

    key = _derive_aead_key(shared_secret, salt)
    aesgcm = AESGCM(key)
    try:
        plaintext = aesgcm.decrypt(nonce, ciphertext, aad_bytes)
    except InvalidTag as exc:  # pragma: no cover - exercised via tests
        raise ValueError("Ciphertext authentication failed") from exc
    finally:
        if isinstance(key, bytes):
            mutable_key = bytearray(key)
            for i in range(len(mutable_key)):
                mutable_key[i] = 0

    return plaintext


def save_encrypted_file(
    payload: EncryptedPayload,
    author: str,
    filename: str = "encrypted_ip.json",
    *,
    metadata: Optional[Dict[str, Any]] = None,
) -> None:
    """Persist an encrypted payload with ISO 8601 metadata."""

    record: Dict[str, Any] = {
        "author": author,
        "encryption_timestamp": datetime.now(timezone.utc).isoformat(),
        "payload": asdict(payload),
    }
    if metadata:
        record["metadata"] = metadata

    with open(filename, "w", encoding="utf-8") as handle:
        json.dump(record, handle, indent=2)
        handle.write("\n")


def load_encrypted_file(path: str | Path) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _write_bytes(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "wb") as handle:
        handle.write(data)


def _read_bytes(path: Path) -> bytes:
    with open(path, "rb") as handle:
        return handle.read()


def _add_generate_keys_subparser(subparsers: Any) -> None:
    parser = subparsers.add_parser("generate-keys", help="Generate a fresh ECC keypair")
    parser.add_argument("--private-out", required=True, help="Path to write the private key PEM")
    parser.add_argument("--public-out", required=True, help="Path to write the public key PEM")
    parser.add_argument(
        "--password",
        help="Optional password to encrypt the private key PEM (use strong passphrases)",
    )


def _add_encrypt_subparser(subparsers: Any) -> None:
    parser = subparsers.add_parser("encrypt", help="Encrypt plaintext for a recipient")
    parser.add_argument("--public-key", required=True, help="Path to the recipient public key PEM")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--plaintext", help="Inline plaintext to encrypt")
    group.add_argument("--infile", help="Path to a UTF-8 text file to encrypt")
    parser.add_argument("--aad", help="Optional associated data to bind to the ciphertext")
    parser.add_argument("--out", required=True, help="Path to write the encrypted JSON payload")


def _add_decrypt_subparser(subparsers: Any) -> None:
    parser = subparsers.add_parser("decrypt", help="Decrypt a payload using the recipient private key")
    parser.add_argument("--private-key", required=True, help="Path to the private key PEM")
    parser.add_argument("--password", help="Password for the encrypted private key, if any")
    parser.add_argument("--input", required=True, help="Path to the encrypted JSON payload")
    parser.add_argument(
        "--aad",
        help="Associated data expected to match encryption (defaults to payload embedded value)",
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Authenticated ECC encryption toolkit")
    subparsers = parser.add_subparsers(dest="command", required=True)
    _add_generate_keys_subparser(subparsers)
    _add_encrypt_subparser(subparsers)
    _add_decrypt_subparser(subparsers)
    return parser


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "generate-keys":
        private_key, public_key = generate_ecc_keys()
        password_bytes = args.password.encode("utf-8") if args.password else None
        private_bytes = serialize_private_key(private_key, password_bytes)
        public_bytes = public_key.public_bytes(
            serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo
        )
        _write_bytes(Path(args.private_out), private_bytes)
        _write_bytes(Path(args.public_out), public_bytes)
        return 0

    if args.command == "encrypt":
        public_bytes = _read_bytes(Path(args.public_key))
        public_key = serialization.load_pem_public_key(public_bytes)
        if args.plaintext is not None:
            plaintext = args.plaintext
        else:
            plaintext = Path(args.infile).read_text(encoding="utf-8")
        payload = encrypt_data(public_key, plaintext, associated_data=args.aad)
        save_encrypted_file(payload, author="unknown", filename=args.out)
        return 0

    if args.command == "decrypt":
        pem_data = _read_bytes(Path(args.private_key))
        password_bytes = args.password.encode("utf-8") if args.password else None
        private_key = load_private_key(pem_data, password_bytes)
        record = load_encrypted_file(args.input)
        payload = EncryptedPayload(**record["payload"])
        plaintext = decrypt_data(private_key, payload, associated_data=args.aad)
        print(plaintext.decode("utf-8"))
        return 0

    parser.error("Unknown command")
    return 1


if __name__ == "__main__":  # pragma: no cover - CLI passthrough
    raise SystemExit(main())