Restructure compliance tools

scripts/compliance/decryption/crypto.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import base64
+import os
+from typing import Any, Dict, Iterable, Optional
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+
+def load_compliance_private_key(key_file: str = "compliance_keypair.txt") -> X25519PrivateKey:
+    if not os.path.exists(key_file):
+        raise FileNotFoundError(f"Compliance key file not found: {key_file}")
+
+    with open(key_file, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    private_key_b64: Optional[str] = None
+    for line in content.split("\n"):
+        line = line.strip()
+        # Look for PRIVATE_KEY= line or base64 lines that are exactly 43 chars (X25519 private key length when base64 encoded)
+        if line.startswith("PRIVATE_KEY="):
+            private_key_b64 = line.split("=", 1)[1].strip()
+            break
+        elif len(line) == 43 and line.endswith("=") and "=" in line:  # Base64 X25519 private key
+            private_key_b64 = line
+            break
+
+    if not private_key_b64:
+        raise ValueError(f"Could not find private key in {key_file}. Expected PRIVATE_KEY= line or 43-character base64 string.")
+
+    private_key_bytes = base64.b64decode(private_key_b64)
+    return X25519PrivateKey.from_private_bytes(private_key_bytes)
+
+
+def _hkdf_32(info: bytes) -> HKDF:
+    return HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=b"\x00" * 16,
+        info=info,
+    )
+
+
+def derive_wrap_key_from_public_key_bytes(public_key_bytes: bytes, context: str) -> bytes:
+    return _hkdf_32(context.encode("utf-8")).derive(public_key_bytes)
+
+
+def derive_compliance_wrap_key(compliance_public_key: X25519PublicKey) -> bytes:
+    return _hkdf_32(b"compliance_wrap_key").derive(compliance_public_key.public_bytes_raw())
+
+
+def decrypt_compliance_mek(
+    wrapped_mek_b64: str,
+    compliance_private_key: X25519PrivateKey,
+    compliance_public_key: X25519PublicKey,
+) -> bytes:
+    wrap_key = derive_compliance_wrap_key(compliance_public_key)
+    wrapped_mek_bytes = base64.b64decode(wrapped_mek_b64)
+    nonce = wrapped_mek_bytes[:12]
+    ciphertext = wrapped_mek_bytes[12:]
+    aesgcm = AESGCM(wrap_key)
+    return aesgcm.decrypt(nonce, ciphertext, None)
+
+
+def decrypt_wrapped_mek_with_public_key(wrapped_mek_b64: str, wrap_public_key_b64: str, wrap_context: str) -> bytes:
+    public_key_bytes = base64.b64decode(wrap_public_key_b64)
+    wrap_key = derive_wrap_key_from_public_key_bytes(public_key_bytes, wrap_context)
+    wrapped_mek_bytes = base64.b64decode(wrapped_mek_b64)
+    nonce = wrapped_mek_bytes[:12]
+    ciphertext = wrapped_mek_bytes[12:]
+    aesgcm = AESGCM(wrap_key)
+    return aesgcm.decrypt(nonce, ciphertext, None)
+
+
+def first_present_key(data: Dict[str, Any], keys: Iterable[str]) -> Optional[str]:
+    for k in keys:
+        v = data.get(k)
+        if v is None:
+            continue
+        if isinstance(v, str) and v.strip() == "":
+            continue
+        return k
+    return None
+
+
+def get_str(data: Dict[str, Any], keys: Iterable[str], label: str) -> str:
+    k = first_present_key(data, keys)
+    if not k:
+        raise ValueError(f"Missing {label}. Expected one of: {', '.join(keys)}")
+    v = data.get(k)
+    if not isinstance(v, str):
+        raise ValueError(f"Invalid {label}: expected string at '{k}', got {type(v).__name__}")
+    return v
+
+
+def decrypt_message(envelope_data: Dict[str, Any], compliance_private_key: X25519PrivateKey, compliance_public_key: X25519PublicKey) -> str:
+    compliance_wrapped_mek = envelope_data.get("compliance_wrapped_mek_b64")
+    if not compliance_wrapped_mek:
+        raise ValueError("Message does not have compliance MEK")
+
+    mek = decrypt_compliance_mek(compliance_wrapped_mek, compliance_private_key, compliance_public_key)
+
+    nonce_b64 = envelope_data["iv_b64"]
+    ciphertext_b64 = envelope_data["ciphertext_b64"]
+
+    nonce = base64.b64decode(nonce_b64)
+    ciphertext = base64.b64decode(ciphertext_b64)
+
+    aesgcm = AESGCM(mek)
+    plaintext = aesgcm.decrypt(nonce, ciphertext, None)
+    return plaintext.decode("utf-8")
+
+
+def decrypt_file_bytes_from_meta(meta: Dict[str, Any], encrypted_bytes: bytes, *, key_file: str = "compliance_keypair.txt") -> bytes:
+    nonce_b64 = get_str(meta, keys=["nonce_b64", "iv_b64", "nonce", "iv"], label="nonce/iv (base64)")
+    nonce = base64.b64decode(nonce_b64)
+
+    mek_key = first_present_key(meta, ["compliance_wrapped_mek_b64", "compliance_wrapped_mek"])
+    if mek_key:
+        compliance_private_key = load_compliance_private_key(key_file=key_file)
+        compliance_public_key = compliance_private_key.public_key()
+        mek = decrypt_compliance_mek(str(meta[mek_key]), compliance_private_key, compliance_public_key)
+    else:
+        wrap_public_key_b64 = get_str(
+            meta,
+            keys=["wrap_public_key_b64", "wrap_public_key", "public_key_b64"],
+            label="wrap public key (base64)",
+        )
+        wrap_context = get_str(meta, keys=["wrap_context"], label="wrap context")
+        wrapped_mek_b64 = get_str(meta, keys=["wrapped_mek_b64", "wrapped_mek"], label="wrapped MEK (base64)")
+        mek = decrypt_wrapped_mek_with_public_key(wrapped_mek_b64, wrap_public_key_b64, wrap_context)
+
+    aesgcm = AESGCM(mek)
+    return aesgcm.decrypt(nonce, encrypted_bytes, None)
+
+
+def derive_auth_secret(username: str, password: str) -> str:
+    """
+    Match frontend `deriveAuthSecret()`:
+    HKDF-SHA256 with:
+        - IKM: UTF-8 password
+        - salt: UTF-8 `fromchat.user:{username}`
+        - info: UTF-8 `auth-secret`
+        - length: 32 bytes
+    Output: base64 string.
+    """
+    salt = f"fromchat.user:{(username or '').strip()}".encode("utf-8")
+    info = b"auth-secret"
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        info=info,
+    )
+    derived = hkdf.derive((password or "").encode("utf-8"))
+    return base64.b64encode(derived).decode("ascii")
+