Integration

"""Send a fax over fxe.

Top-level entry point: ``send_fax(keys_zip, recipient_fax, file_path, from_number)``.
Takes the keys.zip you downloaded after redeeming your invite plus the
recipient fax number + the file to send. Returns the relay's message_id.

Public PyPI deps only:
    pip install httpx cryptography liboqs-python
(liboqs-python also requires the liboqs native library — see
 https://github.com/open-quantum-safe/liboqs.)
"""

from __future__ import annotations

import base64
import hashlib
import io
import json
import os
import secrets
import time
import uuid
import zipfile
from pathlib import Path

import httpx
import oqs  # liboqs-python
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM, ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

SUITE = "fxe-v1-mlkem768-x25519-xchacha20"

# Domain-separation tags. Each is length-prefixed so signatures from one
# domain can't be replayed as another. Format: byte(len) || tag || payload.
DOMAIN_ENVELOPE = b"\x0f" + b"fxe-v1-envelope"  # 15 chars
DOMAIN_BUNDLE   = b"\x0d" + b"fxe-v1-bundle"    # 13 chars

AAD_DOMAIN = b"fxe-v1-aad"
HKDF_INFO_WRAP = b"fxe-v1-wrap"


class CoverageError(Exception):
    """The recipient fax isn't on fxe. Fall back to PSTN."""


def send_fax(
    keys_zip: Path,
    recipient_fax: str,
    file_path: Path,
    from_number: str,
    from_name: str | None = None,
    metadata: dict | None = None,
) -> str:
    """Send `file_path` to `recipient_fax`.

    `from_number` is your caller-ID. It MUST be one of your org's registered
    fax numbers; the relay rejects mismatches with 403.
    """
    creds, bundle_sk = _load_keys(keys_zip)
    relay = creds["relay"]
    api_key = creds["api_key"]

    # 1. Ask the relay which org currently owns this fax number.
    lookup = _request(
        relay, api_key, "GET",
        f"/v1/lookup?fax_number={recipient_fax}",
    )
    if lookup.status_code == 404:
        raise CoverageError(f"{recipient_fax} not covered by fxe")
    lookup.raise_for_status()
    recipient = lookup.json()

    # 2. Encrypt + sign locally. No plaintext or CEK leaves this machine.
    meta = metadata or {"filename": file_path.name, "content_type": "application/pdf"}
    envelope_json, ciphertext = _encrypt_for(
        plaintext=file_path.read_bytes(),
        sender_org_id=creds["org_id"],
        sender_key_bundle_id=creds["bundle_id"],
        sender_bundle_sk=bundle_sk,
        recipient_org_id=recipient["org_id"],
        recipient_bundles=recipient["bundles"],
        from_number=from_number,
        from_name=from_name,
        to_number=recipient_fax,
        metadata=meta,
    )

    # 3. Upload as multipart: signed envelope (JSON) + opaque ciphertext.
    body, content_type = _multipart(envelope_json, ciphertext)
    upload = _request(
        relay, api_key, "POST", "/v1/messages",
        body=body, content_type=content_type,
    )
    upload.raise_for_status()
    return upload.json()["message_id"]


# --- crypto ----------------------------------------------------------------

def _encrypt_for(
    *,
    plaintext: bytes,
    sender_org_id: str,
    sender_key_bundle_id: str,
    sender_bundle_sk: Ed25519PrivateKey,
    recipient_org_id: str,
    recipient_bundles: list[dict],
    from_number: str,
    from_name: str | None,
    to_number: str,
    metadata: dict,
) -> tuple[dict, bytes]:
    """Build a signed envelope + opaque ciphertext for one or more recipient
    bundles of the recipient org (multiple = sender's view of rotation overlap)."""
    message_id = str(uuid.uuid4())
    cek = os.urandom(32)
    nonce = os.urandom(24)  # full 24-byte nonce stored; we use the first 12 today
    aad = _aad(message_id, sender_org_id, recipient_org_id, metadata)
    ct_body = ChaCha20Poly1305(cek).encrypt(nonce[:12], plaintext, aad)
    ciphertext = nonce + ct_body
    ciphertext_hash = hashlib.sha256(ciphertext).digest()

    wrapped_keys = []
    for b in recipient_bundles:
        ml_kem_pub = _b64d(b["ml_kem_pub"])
        x25519_pub = _b64d(b["x25519_pub"])
        # ML-KEM-768 encapsulate against recipient.
        with oqs.KeyEncapsulation("ML-KEM-768") as kem:
            kem_ct, ss_pq = kem.encap_secret(ml_kem_pub)
        # Ephemeral X25519 -> classical shared secret.
        eph_sk = X25519PrivateKey.generate()
        eph_pub = eph_sk.public_key().public_bytes_raw()
        ss_classical = eph_sk.exchange(X25519PublicKey.from_public_bytes(x25519_pub))
        # HKDF binds both shared secrets and the recipient's public keys into
        # a single wrap key — thwarts unknown-key-share attacks.
        info = HKDF_INFO_WRAP + b"|" + eph_pub + b"|" + x25519_pub + b"|" + ml_kem_pub
        wrap_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=info).derive(
            ss_pq + ss_classical
        )
        wnonce = os.urandom(12)
        wrapped_cek = AESGCM(wrap_key).encrypt(wnonce, cek, aad)
        wrapped_keys.append({
            "recipient_key_bundle_id": b["bundle_id"],
            "kem_ct": _b64u(kem_ct),
            "x25519_ephemeral_pub": _b64u(eph_pub),
            "wrap_nonce": _b64u(wnonce),
            "wrapped_cek": _b64u(wrapped_cek),
        })

    created_at = int(time.time())
    # Canonical signing payload — sorted keys, no whitespace, so independent
    # clients produce identical bytes for the same envelope.
    sig_payload = json.dumps(
        {
            "message_id": message_id,
            "suite": SUITE,
            "sender_org_id": sender_org_id,
            "sender_key_bundle_id": sender_key_bundle_id,
            "recipient_org_id": recipient_org_id,
            "from_number": from_number,
            "from_name": from_name,
            "to_number": to_number,
            "ciphertext_sha256": ciphertext_hash.hex(),
            "metadata_sha256": hashlib.sha256(
                json.dumps(metadata, sort_keys=True).encode()
            ).hexdigest(),
            "created_at": created_at,
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    sender_sig = sender_bundle_sk.sign(DOMAIN_ENVELOPE + sig_payload)

    envelope = {
        "message_id": message_id,
        "suite": SUITE,
        "sender_org_id": sender_org_id,
        "sender_key_bundle_id": sender_key_bundle_id,
        "recipient_org_id": recipient_org_id,
        "from_number": from_number,
        "from_name": from_name,
        "to_number": to_number,
        "size": len(ciphertext),
        "ciphertext_sha256": ciphertext_hash.hex(),
        "metadata": metadata,
        "created_at": created_at,
        "sender_sig": _b64u(sender_sig),
        "wrapped_keys": wrapped_keys,
    }
    return envelope, ciphertext


def _aad(message_id: str, sender: str, recipient: str, metadata: dict) -> bytes:
    return (
        AAD_DOMAIN
        + b"|" + message_id.encode()
        + b"|" + sender.encode()
        + b"|" + recipient.encode()
        + b"|" + hashlib.sha256(json.dumps(metadata, sort_keys=True).encode()).digest()
    )


# --- HTTP auth -------------------------------------------------------------

def _request(relay, api_key, method, path, *, body=b"", content_type=None):
    """Make an authenticated request. Auth is a Bearer API key."""
    headers = {"Authorization": f"Bearer {api_key}"}
    if content_type:
        headers["Content-Type"] = content_type
    return httpx.request(method, relay + path, headers=headers, content=body, timeout=60)


# --- helpers ---------------------------------------------------------------

def _load_keys(keys_zip: Path) -> tuple[dict, Ed25519PrivateKey]:
    """Open the keys.zip your /redeem download produced and load the bundle's
    Ed25519 private key — used to sign the per-message envelope. The Bearer
    API key for transport auth is stored as plaintext in credentials.json."""
    with zipfile.ZipFile(keys_zip) as z:
        creds = json.loads(z.read("credentials.json"))
        bundle_seed = z.read("keys/bundle.ed25519.key")
    return creds, Ed25519PrivateKey.from_private_bytes(bundle_seed)


def _b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def _b64d(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def _multipart(envelope: dict, ciphertext: bytes) -> tuple[bytes, str]:
    boundary = "----fxe" + secrets.token_hex(8)
    body = io.BytesIO()
    body.write(f"--{boundary}\r\n".encode())
    body.write(b'Content-Disposition: form-data; name="envelope"\r\n')
    body.write(b"Content-Type: application/json\r\n\r\n")
    body.write(json.dumps(envelope).encode() + b"\r\n")
    body.write(f"--{boundary}\r\n".encode())
    body.write(b'Content-Disposition: form-data; name="ciphertext"; filename="ct.bin"\r\n')
    body.write(b"Content-Type: application/octet-stream\r\n\r\n")
    body.write(ciphertext + b"\r\n")
    body.write(f"--{boundary}--\r\n".encode())
    return body.getvalue(), f"multipart/form-data; boundary={boundary}"


if __name__ == "__main__":
    import sys
    # usage: send.py keys.zip +15559876543 fax.pdf +15551110001 ["Acme Fax"]
    mid = send_fax(
        keys_zip=Path(sys.argv[1]),
        recipient_fax=sys.argv[2],
        file_path=Path(sys.argv[3]),
        from_number=sys.argv[4],
        from_name=sys.argv[5] if len(sys.argv) > 5 else None,
    )
    print(f"sent: {mid}")

"""Receive fxe webhooks (FastAPI).

When a `message.received` event arrives, fetch the envelope, decrypt locally,
then POST a signed receipt back to /v1/messages/{id}/receipt so the relay can
delete the ciphertext.

Auth is one of two modes you picked when you registered the webhook:
  - URL-embedded: the secret is in the URL the relay POSTs to (query, path,
    or userinfo). Pull it out of the request and constant-time compare.
  - Custom header: the relay sends a header you chose with a value you set.
    Constant-time compare against your stored value.

Public PyPI deps only:
    pip install httpx fastapi uvicorn cryptography liboqs-python
"""

from __future__ import annotations

import base64
import hashlib
import json
import secrets as pysecrets
import time
import zipfile
from pathlib import Path

import oqs
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM, ChaCha20Poly1305
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from fastapi import FastAPI, HTTPException, Request

# Reuse the helpers from send.py — sibling file, same dir.
from send import _load_keys, _request, _b64u, _b64d  # noqa: E402

# Domain tags must match those in send.py / the fxe wire protocol.
DOMAIN_RECEIPT = b"\x0e" + b"fxe-v1-receipt"   # 14 chars
DOMAIN_ENVELOPE = b"\x0f" + b"fxe-v1-envelope"  # 15 chars
AAD_DOMAIN = b"fxe-v1-aad"
HKDF_INFO_WRAP = b"fxe-v1-wrap"

app = FastAPI()

KEYS_ZIP = Path("/etc/fxe/keys.zip")

# Pick ONE of these auth modes to match how you registered the webhook.
EXPECTED_HEADER_NAME = "Authorization"
EXPECTED_HEADER_VALUE = b"Bearer your-token"
# OR, for URL-embedded auth:
EXPECTED_URL_TOKEN = "abc123"   # checked against request.query_params["token"]


@app.post("/fxe-webhook")
async def webhook(request: Request):
    body = await request.body()

    if not _auth_ok(request):
        raise HTTPException(401, "bad auth")

    event = json.loads(body)
    event_id = event["event_id"]
    event_type = event["event_type"]

    # Idempotency: if you've processed this event_id before, return 200 and skip.
    if already_handled(event_id):
        return {"ok": True}

    if event_type == "message.received":
        await handle_received(event["data"]["message_id"])

    mark_handled(event_id)
    return {"ok": True}


def _auth_ok(request: Request) -> bool:
    # Custom-header mode — relay sends {EXPECTED_HEADER_NAME: EXPECTED_HEADER_VALUE}.
    got = request.headers.get(EXPECTED_HEADER_NAME, "").encode()
    if pysecrets.compare_digest(got, EXPECTED_HEADER_VALUE):
        return True
    # URL-embedded mode — secret lives in the query string. Uncomment if you
    # registered the URL with ?token=… instead of a custom header.
    # if pysecrets.compare_digest(request.query_params.get("token", ""), EXPECTED_URL_TOKEN):
    #     return True
    return False


async def handle_received(message_id: str) -> None:
    creds, bundle_sk = _load_keys(KEYS_ZIP)
    relay = creds["relay"]
    api_key = creds["api_key"]

    # 1. Pull the envelope + ciphertext (Bearer auth).
    res = _request(relay, api_key, "GET", f"/v1/messages/{message_id}")
    res.raise_for_status()
    detail = res.json()
    envelope = detail["envelope"]
    ciphertext = _b64d(detail["ciphertext"])  # relay returns base64url

    # 2. Decrypt locally with your bundle's KEM private keys (loaded from keys.zip).
    # The envelope ships the sender's bundle Ed25519 pub inline so we can verify
    # `sender_sig` without a second round-trip to /v1/key-bundles.
    sender_signing_pub = _b64d(envelope["sender_signing_pub"])
    plaintext = _decrypt(envelope, ciphertext, creds, sender_signing_pub)
    persist_to_inbox(message_id, plaintext, envelope.get("metadata", {}))

    # 3. Sign a delivery receipt so the relay deletes the ciphertext.
    receipt = _sign_receipt(envelope, creds, bundle_sk)
    _request(
        relay, api_key, "POST",
        f"/v1/messages/{message_id}/receipt",
        body=json.dumps(receipt).encode(),
        content_type="application/json",
    )


# --- crypto ----------------------------------------------------------------

def _decrypt(
    envelope: dict, ciphertext: bytes, creds: dict, sender_signing_pub: bytes,
) -> bytes:
    """Verify the sender's envelope signature, then unwrap our CEK and
    AEAD-decrypt the ciphertext. Raises on any tamper."""
    # 1. Reconstruct the canonical signing payload and verify.
    sig_payload = json.dumps(
        {
            "message_id": envelope["message_id"],
            "suite": envelope["suite"],
            "sender_org_id": envelope["sender_org_id"],
            "sender_key_bundle_id": envelope["sender_key_bundle_id"],
            "recipient_org_id": envelope["recipient_org_id"],
            "from_number": envelope["from_number"],
            "from_name": envelope.get("from_name"),
            "to_number": envelope["to_number"],
            "ciphertext_sha256": envelope["ciphertext_sha256"],
            "metadata_sha256": hashlib.sha256(
                json.dumps(envelope.get("metadata", {}), sort_keys=True).encode()
            ).hexdigest(),
            "created_at": envelope["created_at"],
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    try:
        Ed25519PublicKey.from_public_bytes(sender_signing_pub).verify(
            _b64d(envelope["sender_sig"]),
            DOMAIN_ENVELOPE + sig_payload,
        )
    except InvalidSignature:
        raise ValueError("envelope signature invalid")

    # 2. Verify the ciphertext hasn't been mangled in transit.
    if hashlib.sha256(ciphertext).hexdigest() != envelope["ciphertext_sha256"]:
        raise ValueError("ciphertext hash mismatch")

    # 3. Find the wrapped CEK addressed to *our* bundle and unwrap it.
    our_bundle = creds["bundle_id"]
    wk = next(
        (w for w in envelope["wrapped_keys"] if w["recipient_key_bundle_id"] == our_bundle),
        None,
    )
    if wk is None:
        raise KeyError(f"no wrapped key for our bundle {our_bundle}")

    ml_kem_sk = _read_secret("keys/kem.mlkem768.key")
    ml_kem_pub = _read_secret("keys/kem.mlkem768.pub")
    x25519_sk = _read_secret("keys/kem.x25519.key")
    x25519_pub = _read_secret("keys/kem.x25519.pub")

    with oqs.KeyEncapsulation("ML-KEM-768", secret_key=ml_kem_sk) as kem:
        ss_pq = kem.decap_secret(_b64d(wk["kem_ct"]))
    eph_pub = _b64d(wk["x25519_ephemeral_pub"])
    sk = X25519PrivateKey.from_private_bytes(x25519_sk)
    ss_classical = sk.exchange(X25519PublicKey.from_public_bytes(eph_pub))

    info = HKDF_INFO_WRAP + b"|" + eph_pub + b"|" + x25519_pub + b"|" + ml_kem_pub
    wrap_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=info,
    ).derive(ss_pq + ss_classical)

    aad = _aad(
        envelope["message_id"], envelope["sender_org_id"],
        envelope["recipient_org_id"], envelope.get("metadata", {}),
    )
    cek = AESGCM(wrap_key).decrypt(_b64d(wk["wrap_nonce"]), _b64d(wk["wrapped_cek"]), aad)

    # 4. Finally, AEAD-decrypt the body.
    nonce = ciphertext[:24]
    body = ciphertext[24:]
    return ChaCha20Poly1305(cek).decrypt(nonce[:12], body, aad)


def _sign_receipt(envelope: dict, creds: dict, bundle_sk: Ed25519PrivateKey) -> dict:
    """Sign a delivery receipt the relay verifies before purging the blob."""
    now = int(time.time())
    payload = json.dumps(
        {
            "ciphertext_sha256": envelope["ciphertext_sha256"],
            "message_id": envelope["message_id"],
            "recipient_bundle_id": creds["bundle_id"],
            "recipient_org_id": creds["org_id"],
            "signed_at": now,
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    sig = bundle_sk.sign(DOMAIN_RECEIPT + payload)
    return {
        "message_id": envelope["message_id"],
        "ciphertext_sha256": envelope["ciphertext_sha256"],
        "recipient_bundle_id": creds["bundle_id"],
        "signed_at": now,
        "recipient_sig": _b64u(sig),
    }


def _aad(message_id: str, sender: str, recipient: str, metadata: dict) -> bytes:
    return (
        AAD_DOMAIN
        + b"|" + message_id.encode()
        + b"|" + sender.encode()
        + b"|" + recipient.encode()
        + b"|" + hashlib.sha256(json.dumps(metadata, sort_keys=True).encode()).digest()
    )


def _read_secret(name: str) -> bytes:
    """Pull a raw key file out of the keys.zip without re-extracting everything."""
    with zipfile.ZipFile(KEYS_ZIP) as z:
        return z.read(name)


# --- placeholders you implement -------------------------------------------

def already_handled(event_id: str) -> bool: ...
def mark_handled(event_id: str) -> None: ...
def persist_to_inbox(message_id: str, plaintext: bytes, metadata: dict) -> None: ...