Add AES-GCM file encryption (encrypt_file / decrypt_file)

JE-Chen · JE-Chen · commit b50c066c1093 · 2026-04-21T22:36:42.000+08:00
diff --git a/automation_file/__init__.py b/automation_file/__init__.py
@@ -29,6 +29,13 @@
 )
 from automation_file.core.config import AutomationConfig, ConfigException
 from automation_file.core.config_watcher import ConfigWatcher
+from automation_file.core.crypto import (
+    CryptoException,
+    decrypt_file,
+    encrypt_file,
+    generate_key,
+    key_from_password,
+)
 from automation_file.core.dag_executor import execute_action_dag
 from automation_file.core.fim import IntegrityMonitor
 from automation_file.core.json_store import read_action_json, write_action_json
@@ -325,6 +332,11 @@ def __getattr__(name: str) -> Any:
     "AuditException",
     "AuditLog",
     "IntegrityMonitor",
+    "CryptoException",
+    "encrypt_file",
+    "decrypt_file",
+    "generate_key",
+    "key_from_password",
     # Triggers
     "FileWatcher",
     "TriggerManager",
diff --git a/automation_file/core/action_registry.py b/automation_file/core/action_registry.py
@@ -153,7 +153,7 @@ def _http_commands() -> dict[str, Command]:
 
 
 def _utils_commands() -> dict[str, Command]:
-    from automation_file.core import checksum, manifest
+    from automation_file.core import checksum, crypto, manifest
     from automation_file.remote import cross_backend
     from automation_file.utils import deduplicate, fast_find, grep, rotate
 
@@ -168,6 +168,8 @@ def _utils_commands() -> dict[str, Command]:
         "FA_grep": grep.grep_files,
         "FA_rotate_backups": rotate.rotate_backups,
         "FA_copy_between": cross_backend.copy_between,
+        "FA_encrypt_file": crypto.encrypt_file,
+        "FA_decrypt_file": crypto.decrypt_file,
     }
 
 
diff --git a/automation_file/core/crypto.py b/automation_file/core/crypto.py
@@ -0,0 +1,174 @@
+"""AES-256-GCM file encryption helpers.
+
+``encrypt_file(source, target, key)`` writes a self-describing envelope::
+
+    magic    = b"FA-AESG"     7 bytes
+    version  = 0x01           1 byte
+    flags    = 0x00           1 byte (reserved)
+    aad_len  = uint32 BE      4 bytes
+    nonce    = 12 bytes
+    aad      = <aad_len>
+    ciphertext + tag          (rest — GCM tag is the trailing 16 bytes)
+
+``decrypt_file`` reads the same format, verifies the tag, and writes the
+plaintext to ``target``. Tampering (bit flips anywhere in the envelope
+except ``aad_len``) surfaces as :class:`CryptoException`.
+
+GCM has a hard plaintext limit of roughly 64 GiB per ``(key, nonce)``
+pair; since each encrypt generates a fresh nonce, the practical cap is
+per-file and is much larger than typical automation payloads. For files
+approaching that size, split before calling ``encrypt_file``.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from cryptography.exceptions import InvalidTag
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+from automation_file.exceptions import FileAutomationException
+from automation_file.logging_config import file_automation_logger
+
+_MAGIC = b"FA-AESG"
+_VERSION = 0x01
+_NONCE_SIZE = 12
+_HEADER_FIXED_SIZE = len(_MAGIC) + 2 + 4  # magic + version + flags + aad_len
+_VALID_KEY_SIZES = frozenset({16, 24, 32})
+_DEFAULT_PBKDF2_ITERATIONS = 200_000
+
+
+class CryptoException(FileAutomationException):
+    """Raised when encryption / decryption fails (including on tamper)."""
+
+
+def generate_key(*, bits: int = 256) -> bytes:
+    """Return cryptographically random bytes suitable for AES-GCM."""
+    if bits not in (128, 192, 256):
+        raise CryptoException(f"bits must be 128 / 192 / 256, got {bits}")
+    return os.urandom(bits // 8)
+
+
+def key_from_password(
+    password: str,
+    salt: bytes,
+    *,
+    iterations: int = _DEFAULT_PBKDF2_ITERATIONS,
+    bits: int = 256,
+) -> bytes:
+    """Derive a symmetric key from ``password`` via PBKDF2-HMAC-SHA256."""
+    if not password:
+        raise CryptoException("password must be non-empty")
+    if len(salt) < 16:
+        raise CryptoException("salt must be at least 16 bytes")
+    if bits not in (128, 192, 256):
+        raise CryptoException(f"bits must be 128 / 192 / 256, got {bits}")
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=bits // 8,
+        salt=salt,
+        iterations=iterations,
+    )
+    return kdf.derive(password.encode("utf-8"))
+
+
+def encrypt_file(
+    source: str | os.PathLike[str],
+    target: str | os.PathLike[str],
+    key: bytes,
+    *,
+    associated_data: bytes = b"",
+) -> dict[str, int]:
+    """Encrypt ``source`` to ``target`` under AES-GCM. Returns a size summary."""
+    _validate_key(key)
+    if not isinstance(associated_data, (bytes, bytearray)):
+        raise CryptoException("associated_data must be bytes")
+    src = Path(source)
+    if not src.is_file():
+        raise CryptoException(f"source file not found: {src}")
+
+    plaintext = src.read_bytes()
+    nonce = os.urandom(_NONCE_SIZE)
+    aesgcm = AESGCM(bytes(key))
+    ciphertext = aesgcm.encrypt(nonce, plaintext, bytes(associated_data) or None)
+
+    envelope = _build_header(associated_data, nonce) + ciphertext
+    dst = Path(target)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    dst.write_bytes(envelope)
+    file_automation_logger.info(
+        "encrypt_file: %s -> %s (%d -> %d bytes)",
+        src,
+        dst,
+        len(plaintext),
+        len(envelope),
+    )
+    return {"plaintext_bytes": len(plaintext), "ciphertext_bytes": len(envelope)}
+
+
+def decrypt_file(
+    source: str | os.PathLike[str],
+    target: str | os.PathLike[str],
+    key: bytes,
+) -> dict[str, int]:
+    """Decrypt ``source`` to ``target``. Raises on invalid tag / header."""
+    _validate_key(key)
+    src = Path(source)
+    if not src.is_file():
+        raise CryptoException(f"source file not found: {src}")
+    envelope = src.read_bytes()
+    nonce, aad, ciphertext = _parse_envelope(envelope)
+    aesgcm = AESGCM(bytes(key))
+    try:
+        plaintext = aesgcm.decrypt(nonce, ciphertext, aad or None)
+    except InvalidTag as err:
+        raise CryptoException("authentication failed: wrong key or tampered data") from err
+
+    dst = Path(target)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    dst.write_bytes(plaintext)
+    file_automation_logger.info(
+        "decrypt_file: %s -> %s (%d -> %d bytes)",
+        src,
+        dst,
+        len(envelope),
+        len(plaintext),
+    )
+    return {"ciphertext_bytes": len(envelope), "plaintext_bytes": len(plaintext)}
+
+
+def _validate_key(key: bytes) -> None:
+    if not isinstance(key, (bytes, bytearray)):
+        raise CryptoException("key must be bytes")
+    if len(key) not in _VALID_KEY_SIZES:
+        raise CryptoException(
+            f"key length must be 16 / 24 / 32 bytes, got {len(key)}",
+        )
+
+
+def _build_header(aad: bytes, nonce: bytes) -> bytes:
+    aad_len = len(aad).to_bytes(4, "big")
+    return _MAGIC + bytes([_VERSION, 0x00]) + aad_len + nonce + bytes(aad)
+
+
+def _parse_envelope(envelope: bytes) -> tuple[bytes, bytes, bytes]:
+    if len(envelope) < _HEADER_FIXED_SIZE + _NONCE_SIZE + 16:
+        raise CryptoException("ciphertext envelope is shorter than the fixed header")
+    if not envelope.startswith(_MAGIC):
+        raise CryptoException("not an AES-GCM envelope (bad magic)")
+    version = envelope[len(_MAGIC)]
+    if version != _VERSION:
+        raise CryptoException(f"unsupported envelope version {version}")
+    aad_len = int.from_bytes(envelope[_HEADER_FIXED_SIZE - 4 : _HEADER_FIXED_SIZE], "big")
+    nonce_start = _HEADER_FIXED_SIZE
+    nonce_end = nonce_start + _NONCE_SIZE
+    aad_end = nonce_end + aad_len
+    if aad_end > len(envelope):
+        raise CryptoException("envelope truncated before aad end")
+    nonce = envelope[nonce_start:nonce_end]
+    aad = envelope[nonce_end:aad_end]
+    ciphertext = envelope[aad_end:]
+    return nonce, aad, ciphertext
diff --git a/dev.toml b/dev.toml
@@ -25,6 +25,7 @@ dependencies = [
     "paramiko>=3.4.0",
     "PySide6>=6.6.0",
     "watchdog>=4.0.0",
+    "cryptography>=42.0.0",
     "tomli>=2.0.1; python_version<\"3.11\""
 ]
 classifiers = [
diff --git a/stable.toml b/stable.toml
@@ -25,6 +25,7 @@ dependencies = [
     "paramiko>=3.4.0",
     "PySide6>=6.6.0",
     "watchdog>=4.0.0",
+    "cryptography>=42.0.0",
     "tomli>=2.0.1; python_version<\"3.11\""
 ]
 classifiers = [
diff --git a/tests/test_crypto.py b/tests/test_crypto.py
@@ -0,0 +1,138 @@
+"""Tests for AES-GCM file encryption helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from automation_file import (
+    CryptoException,
+    build_default_registry,
+    decrypt_file,
+    encrypt_file,
+    generate_key,
+    key_from_password,
+)
+
+
+def test_generate_key_default_length() -> None:
+    assert len(generate_key()) == 32
+
+
+def test_generate_key_rejects_bad_size() -> None:
+    with pytest.raises(CryptoException):
+        generate_key(bits=111)
+
+
+def test_round_trip_preserves_plaintext(tmp_path: Path) -> None:
+    key = generate_key()
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"top-secret payload\n")
+    enc = tmp_path / "cipher.bin"
+    dec = tmp_path / "restored.bin"
+    summary = encrypt_file(source, enc, key)
+    assert summary["plaintext_bytes"] == source.stat().st_size
+    decrypt_file(enc, dec, key)
+    assert dec.read_bytes() == source.read_bytes()
+
+
+def test_ciphertext_differs_between_calls(tmp_path: Path) -> None:
+    key = generate_key()
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"same data")
+    enc_a = tmp_path / "a.bin"
+    enc_b = tmp_path / "b.bin"
+    encrypt_file(source, enc_a, key)
+    encrypt_file(source, enc_b, key)
+    assert enc_a.read_bytes() != enc_b.read_bytes()
+
+
+def test_wrong_key_fails(tmp_path: Path) -> None:
+    good = generate_key()
+    bad = generate_key()
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"x" * 128)
+    enc = tmp_path / "cipher.bin"
+    encrypt_file(source, enc, good)
+    with pytest.raises(CryptoException, match="authentication"):
+        decrypt_file(enc, tmp_path / "dec.bin", bad)
+
+
+def test_tampered_ciphertext_fails(tmp_path: Path) -> None:
+    key = generate_key()
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"x" * 128)
+    enc = tmp_path / "cipher.bin"
+    encrypt_file(source, enc, key)
+    contents = bytearray(enc.read_bytes())
+    contents[-1] ^= 0x01
+    enc.write_bytes(bytes(contents))
+    with pytest.raises(CryptoException, match="authentication"):
+        decrypt_file(enc, tmp_path / "dec.bin", key)
+
+
+def test_associated_data_roundtrip(tmp_path: Path) -> None:
+    key = generate_key()
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"hello")
+    enc = tmp_path / "cipher.bin"
+    encrypt_file(source, enc, key, associated_data=b"file=plain.bin")
+    dec = tmp_path / "dec.bin"
+    decrypt_file(enc, dec, key)
+    assert dec.read_bytes() == b"hello"
+
+
+def test_invalid_key_size(tmp_path: Path) -> None:
+    source = tmp_path / "plain.bin"
+    source.write_bytes(b"x")
+    with pytest.raises(CryptoException, match="key length"):
+        encrypt_file(source, tmp_path / "out", b"\x00" * 10)
+
+
+def test_missing_source_raises(tmp_path: Path) -> None:
+    with pytest.raises(CryptoException, match="source file"):
+        encrypt_file(tmp_path / "absent", tmp_path / "out", generate_key())
+
+
+def test_bad_magic_rejected(tmp_path: Path) -> None:
+    garbage = tmp_path / "junk.bin"
+    garbage.write_bytes(b"NOT-A-REAL-ENVELOPE" + b"\x00" * 64)
+    with pytest.raises(CryptoException, match="magic"):
+        decrypt_file(garbage, tmp_path / "out.bin", generate_key())
+
+
+def test_key_from_password_deterministic() -> None:
+    salt = b"\x00" * 16
+    key_a = key_from_password("passphrase", salt, iterations=1_000)
+    key_b = key_from_password("passphrase", salt, iterations=1_000)
+    assert key_a == key_b
+    assert len(key_a) == 32
+
+
+def test_key_from_password_requires_nonempty_password() -> None:
+    with pytest.raises(CryptoException, match="non-empty"):
+        key_from_password("", b"\x00" * 16)
+
+
+def test_key_from_password_rejects_short_salt() -> None:
+    with pytest.raises(CryptoException, match="salt"):
+        key_from_password("pw", b"\x00" * 8)
+
+
+def test_key_from_password_round_trip(tmp_path: Path) -> None:
+    salt = b"sixteen-byte-saltX"
+    key = key_from_password("strong-password", salt, iterations=1_000)
+    source = tmp_path / "plain.txt"
+    source.write_text("hello", encoding="utf-8")
+    enc = tmp_path / "cipher.bin"
+    encrypt_file(source, enc, key)
+    dec = tmp_path / "dec.txt"
+    decrypt_file(enc, dec, key)
+    assert dec.read_text(encoding="utf-8") == "hello"
+
+
+def test_encrypt_decrypt_actions_registered() -> None:
+    registry = build_default_registry()
+    assert "FA_encrypt_file" in registry
+    assert "FA_decrypt_file" in registry

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ dependencies = [`
`25`	`25`	`"paramiko>=3.4.0",`
`26`	`26`	`"PySide6>=6.6.0",`
`27`	`27`	`"watchdog>=4.0.0",`
	`28`	`+ "cryptography>=42.0.0",`
`28`	`29`	`"tomli>=2.0.1; python_version<\"3.11\""`
`29`	`30`	`]`
`30`	`31`	`classifiers = [`