dove-core 0.1.1

"""dove access gate — the full tier's policy enforcer and page server.

Fronted by a Lambda Function URL. Share (download) routes:

  GET /d/<id>/<name>  → serve the decryptor page (no decrement; unfurler-safe)
  GET /meta/<id>      → the share's policy as JSON (free; for the page to show
                        expiry / downloads-left and decide small-vs-large)
  GET /dl/<id>        → check + atomically decrement the download budget, then
                        302 to a short-lived presigned S3 URL

File-request (upload) routes — a request row carries kind="request":

  GET  /r/<id>        → serve the upload page (no side effect; unfurler-safe)
  GET  /rmeta/<id>    → the request's status as JSON (waiting/received/failed)
  GET  /up/<id>       → PIN-gated, rate-limited handout of a presigned S3 POST
                        policy (content-length-range capped); no budget spent
  POST /done/<id>     → finalize: confirm the object, record it, spend the budget

/verify is shared by both and is kind-aware. The gate never sees the decryption
key — that rides the URL fragment, which browsers and HTTP clients never send.
Environment: BUCKET, TABLE.
"""

import base64
import hashlib
import hmac
import json
import os
import time
from pathlib import Path

import boto3

BUCKET = os.environ["BUCKET"]
TABLE = os.environ["TABLE"]

# Wrong-PIN guesses allowed before the share locks. Small, because the gate
# checks online — a handful of tries against a 6-digit PIN is negligible odds,
# and locking makes brute force impossible rather than merely slow.
MAX_PIN_ATTEMPTS = 5

# Presigned-upload handouts a file-request will issue before it rate-limits. The
# budget of *successful* uploads (uploads_remaining) is spent only at /done, so
# this ceiling is the anti-hammer on the handout itself — a failed upload can be
# retried up to this many times without burning the budget.
MAX_UPLOAD_ATTEMPTS = 10

# The content-length-range ceiling baked into every presigned POST policy, so S3
# itself rejects an over-size upload at PUT time (the gate never sees the bytes).
MAX_UPLOAD = 512 * 1024 * 1024  # 512 MiB

# CORS for the JSON the browser upload page fetches cross-origin from the gate.
_CORS = {"access-control-allow-origin": "*"}

_ddb = boto3.client("dynamodb")
_s3 = boto3.client("s3")
_ssm = boto3.client("ssm")


def _load_gate_secret():
    """The HMAC key for share ids: from SSM SecureString (GATE_SECRET_PARAM) —
    read once at cold start, never in the function config. Falls back to a direct
    GATE_SECRET env var for legacy deploys. None → id verification is skipped."""
    name = os.environ.get("GATE_SECRET_PARAM")
    if name:
        try:
            value = _ssm.get_parameter(Name=name, WithDecryption=True)["Parameter"]["Value"]
            return bytes.fromhex(value)
        except Exception:  # noqa: BLE001
            return None
    raw = os.environ.get("GATE_SECRET")
    return bytes.fromhex(raw) if raw else None


GATE_SECRET = _load_gate_secret()
PAGE = (Path(__file__).parent / "share.html").read_text()
REQUEST_PAGE = (Path(__file__).parent / "request.html").read_text()
# The link-preview image messaging apps show when a share link is pasted. Generic
# and branded — it can't reveal the filename (that's E2E), which is the point.
OG_PNG = (Path(__file__).parent / "og.png").read_bytes()
# A distinct card for file *requests*, so a request link's unfurl reads "requested",
# not "shared" — the request page points its og:image here.
OG_REQUEST_PNG = (Path(__file__).parent / "og-request.png").read_bytes()


def _png(data):
    """A cacheable image/png response for the unfurl link-preview cards."""
    return {
        "statusCode": 200,
        "headers": {"content-type": "image/png", "cache-control": "public, max-age=86400"},
        "body": base64.b64encode(data).decode(),
        "isBase64Encoded": True,
    }


def _resp(status, body="", content_type="text/plain; charset=utf-8", extra=None):
    headers = {"content-type": content_type}
    if extra:
        headers.update(extra)
    return {"statusCode": status, "headers": headers, "body": body}


def _origin(event):
    """This gate's own reachable origin, for absolute og:image / og:url."""
    dom = event.get("requestContext", {}).get("domainName", "")
    return f"https://{dom}" if dom else ""


def _valid_id(share_id):
    """True if the id carries a MAC this gate minted: hex(nonce(8) ‖
    HMAC-SHA256(GATE_SECRET, nonce)[:8]). Constant-time; no DB access."""
    if GATE_SECRET is None:
        return True  # verification not configured — don't lock the gate
    if len(share_id) != 32:
        return False
    try:
        raw = bytes.fromhex(share_id)
    except ValueError:
        return False
    nonce, mac = raw[:8], raw[8:]
    expected = hmac.new(GATE_SECRET, nonce, hashlib.sha256).digest()[:8]
    return hmac.compare_digest(mac, expected)


def handler(event, _context):
    raw = event.get("rawPath", "")
    if raw == "/og.png":
        return _png(OG_PNG)  # share-link unfurl (fetched by unfurlers via og:image)
    if raw == "/og-request.png":
        return _png(OG_REQUEST_PNG)  # request-link unfurl — distinct "requested" card
    parts = [p for p in raw.split("/") if p]
    if len(parts) < 2:
        return _resp(404, "not a dove share link")
    route, share_id = parts[0], parts[1]

    # Verify the id's MAC before any DynamoDB/S3 touch — a forged or random id is
    # rejected here, for the cost of a hash. (The pages /d and /r are served
    # regardless; they're static and cheap, and their /meta|/rmeta call is where
    # the real work is gated.)
    if route in ("meta", "verify", "dl", "rmeta", "up", "done") and not _valid_id(share_id):
        return _resp(403, "not a dove share link")

    if route == "d":
        # The decryptor page. No decrement — opening a link (or an unfurler
        # previewing it) never spends a download. Inject this gate's origin so the
        # og:image / og:url are absolute (unfurlers require it).
        page = PAGE.replace("__OGBASE__", _origin(event))
        return _resp(200, page, "text/html; charset=utf-8")

    if route == "r":
        # The upload page for a file-request. Static and NOT MAC-gated (like /d);
        # its /rmeta call is where the request is actually resolved. Same origin
        # injection as /d so its og:image / og:url are absolute.
        page = REQUEST_PAGE.replace("__OGBASE__", _origin(event))
        return _resp(200, page, "text/html; charset=utf-8")

    if route == "meta":
        return _meta(share_id)

    if route == "verify":
        # PIN pre-check for the browser's two-step flow: verify + rate-limit
        # WITHOUT spending a download. The download happens on the explicit click.
        params = event.get("queryStringParameters") or {}
        return _verify(share_id, params.get("pin"))

    if route == "dl":
        params = event.get("queryStringParameters") or {}
        return _download(share_id, params.get("pin"))

    if route == "rmeta":
        return _rmeta(share_id)

    if route == "up":
        params = event.get("queryStringParameters") or {}
        return _up(share_id, params.get("pin"))

    if route == "done":
        params = event.get("queryStringParameters") or {}
        body = event.get("body") or ""
        if event.get("isBase64Encoded"):
            try:
                body = base64.b64decode(body).decode()
            except Exception:  # noqa: BLE001
                body = ""
        return _done(share_id, params.get("pin"), body)

    return _resp(404, "not a dove share link")


def _meta(share_id):
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": share_id}}).get("Item")
    if not item:
        return _resp(404, json.dumps({"error": "not found"}), "application/json")
    # Size is stored on the item (no per-request HeadObject). The filename is NOT
    # here — it's end-to-end encrypted in the link's fragment, which the gate
    # never sees. The page decrypts it client-side.
    size = int(item.get("size", {}).get("N", "0"))
    if not size:  # older shares written before size was stored
        try:
            size = _s3.head_object(Bucket=BUCKET, Key=item["s3_key"]["S"])["ContentLength"]
        except Exception:  # noqa: BLE001
            pass
    pin_required = "pin_hash" in item
    locked = pin_required and int(item.get("pin_attempts", {}).get("N", "0")) >= MAX_PIN_ATTEMPTS
    body = json.dumps(
        {
            "downloads_remaining": int(item["downloads_remaining"]["N"]),
            "downloads_total": int(item.get("downloads_total", {}).get("N", "0")),
            "expires_at": int(item["expires_at"]["N"]),
            "size": size,
            # Opaque encrypted blob (filename + trust); the client decrypts it with
            # the fragment secret. The gate can't read it.
            "meta": item.get("meta", {}).get("S", ""),
            "pin_required": pin_required,
            "locked": locked,
        }
    )
    return _resp(200, body, "application/json", {"access-control-allow-origin": "*"})


def _pin_gate(share_id, item, pin):
    """Verify a PIN-locked share. Returns an error response to send back, or None
    to allow the download. A wrong guess is counted; enough of them lock it."""
    pin_hash = item.get("pin_hash", {}).get("S")
    if not pin_hash:
        return None  # not PIN-locked
    attempts = int(item.get("pin_attempts", {}).get("N", "0"))
    if attempts >= MAX_PIN_ATTEMPTS:
        return _resp(423, json.dumps({"error": "locked"}), "application/json")
    if not pin:
        return _resp(401, json.dumps({"error": "pin required"}), "application/json")
    if hashlib.sha256(f"{share_id}:{pin}".encode()).hexdigest() == pin_hash:
        return None  # correct — allow through
    # Wrong. Count it atomically; the same op locks the share at the ceiling.
    try:
        res = _ddb.update_item(
            TableName=TABLE,
            Key={"id": {"S": share_id}},
            UpdateExpression="SET pin_attempts = if_not_exists(pin_attempts, :z) + :one",
            ConditionExpression="attribute_not_exists(pin_attempts) OR pin_attempts < :max",
            ExpressionAttributeValues={
                ":one": {"N": "1"},
                ":z": {"N": "0"},
                ":max": {"N": str(MAX_PIN_ATTEMPTS)},
            },
            ReturnValues="ALL_NEW",
        )
        remaining = max(0, MAX_PIN_ATTEMPTS - int(res["Attributes"]["pin_attempts"]["N"]))
    except _ddb.exceptions.ConditionalCheckFailedException:
        return _resp(423, json.dumps({"error": "locked"}), "application/json")
    status = 423 if remaining == 0 else 401
    error = "locked" if remaining == 0 else "wrong pin"
    return _resp(
        status,
        json.dumps({"error": error, "attempts_remaining": remaining}),
        "application/json",
    )


def _verify(share_id, pin):
    now = int(time.time())
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": share_id}}).get("Item")
    # Kind-aware liveness. A request row is live while it has upload budget, is
    # unexpired, and hasn't already received a file; a share row keeps its exact
    # original test (downloads_remaining), byte-for-byte.
    if item and item.get("kind", {}).get("S") == "request":
        if (
            int(item["expires_at"]["N"]) <= now
            or int(item.get("uploads_remaining", {}).get("N", "0")) <= 0
            or "s3_key" in item
        ):
            return _resp(410, json.dumps({"error": "gone"}), "application/json")
    elif not item or int(item["expires_at"]["N"]) <= now or int(item["downloads_remaining"]["N"]) <= 0:
        return _resp(410, json.dumps({"error": "gone"}), "application/json")
    # Same PIN gate as the download (verify + rate-limit + lock), but no decrement.
    gate = _pin_gate(share_id, item, pin)
    if gate is not None:
        return gate
    return _resp(200, json.dumps({"ok": True}), "application/json")


def _download(share_id, pin):
    now = int(time.time())
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": share_id}}).get("Item")
    if not item:
        return _resp(410, "this share has expired or reached its download limit")
    if int(item["expires_at"]["N"]) <= now or int(item["downloads_remaining"]["N"]) <= 0:
        return _resp(410, "this share has expired or reached its download limit")

    # Second factor: verify the PIN (if any) before spending a download.
    gate = _pin_gate(share_id, item, pin)
    if gate is not None:
        return gate

    try:
        # Atomic: decrement only if the share exists, has budget, and is unexpired.
        result = _ddb.update_item(
            TableName=TABLE,
            Key={"id": {"S": share_id}},
            UpdateExpression="SET downloads_remaining = downloads_remaining - :one",
            ConditionExpression=(
                "attribute_exists(id) "
                "AND downloads_remaining > :zero "
                "AND expires_at > :now"
            ),
            ExpressionAttributeValues={
                ":one": {"N": "1"},
                ":zero": {"N": "0"},
                ":now": {"N": str(now)},
            },
            ReturnValues="ALL_NEW",
        )
    except _ddb.exceptions.ConditionalCheckFailedException:
        return _resp(410, "this share has expired or reached its download limit")
    except Exception:  # noqa: BLE001 - never leak internals to a downloader
        return _resp(500, "the gate hit an error")

    s3_key = result["Attributes"]["s3_key"]["S"]
    # The presign window governs when the download must *start* (and leaves room
    # to resume a dropped transfer) — NOT how long it may run. S3 checks expiry
    # at request time; once the GET is accepted it streams the whole object even
    # if the window passes, so huge multi-hour downloads are fine as long as they
    # begin within the window (clients follow the 302 immediately). 15 minutes
    # gives resume headroom, well within the Lambda role's credential lifetime.
    presigned = _s3.generate_presigned_url(
        "get_object",
        Params={"Bucket": BUCKET, "Key": s3_key},
        ExpiresIn=900,
    )
    return _resp(302, "", extra={"location": presigned})


# --- file-request routes ---------------------------------------------------
# A request row is a share row's mirror image: instead of a budget of downloads
# it holds a budget of *uploads*. The same id MAC, the same PIN gate, and the
# same atomic conditional-update idiom carry over unchanged.


def _rmeta(request_id):
    """The request's status for the upload page and `dove requests`. Free (no
    decrement); CORS-open so the browser page can fetch it cross-origin."""
    now = int(time.time())
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": request_id}}).get("Item")
    if not item:
        return _resp(404, json.dumps({"error": "not found"}), "application/json", _CORS)

    pin_required = "pin_hash" in item
    pin_attempts = int(item.get("pin_attempts", {}).get("N", "0"))
    upload_attempts = int(item.get("upload_attempts", {}).get("N", "0"))
    expires_at = int(item["expires_at"]["N"])
    received = "s3_key" in item
    locked = pin_attempts >= MAX_PIN_ATTEMPTS

    # Status, in strict priority order (a received file wins over any failure;
    # failures rank locked > rate-limited > expired).
    reason = None
    if received:
        status = "received"
    elif pin_attempts >= MAX_PIN_ATTEMPTS:
        status, reason = "failed", "locked"
    elif upload_attempts >= MAX_UPLOAD_ATTEMPTS:
        status, reason = "failed", "rate-limited"
    elif expires_at <= now:
        status, reason = "failed", "expired"
    else:
        status = "waiting"

    body = {
        # Opaque encrypted trust blob (from/message/desc); the client decrypts it
        # with the fragment secret. The gate can't read it.
        "meta": item.get("meta", {}).get("S", ""),
        "expires_at": expires_at,
        "uploads_remaining": int(item.get("uploads_remaining", {}).get("N", "0")),
        "pin_required": pin_required,
        "locked": locked,
        "status": status,
        "reason": reason,
    }
    if received:
        # The filename blob the uploader sealed (encrypted client-side) and the
        # confirmed object size, so the requester's page can name and size it.
        body["name_meta"] = item.get("upload_meta", {}).get("S", "")
        body["size"] = int(item.get("size", {}).get("N", "0"))
    return _resp(200, json.dumps(body), "application/json", _CORS)


def _up(request_id, pin):
    """Hand out a short-lived presigned S3 POST policy for the uploader. PIN-gated
    and rate-limited; the upload budget is NOT spent here (only at /done), so a
    failed upload can be retried up to MAX_UPLOAD_ATTEMPTS times."""
    now = int(time.time())
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": request_id}}).get("Item")
    if not item or item.get("kind", {}).get("S") != "request":
        return _resp(404, json.dumps({"error": "not found"}), "application/json", _CORS)
    if (
        int(item["expires_at"]["N"]) <= now
        or "s3_key" in item  # already received
        or int(item.get("uploads_remaining", {}).get("N", "0")) <= 0
    ):
        return _resp(410, json.dumps({"error": "gone"}), "application/json", _CORS)

    # Second factor: verify the PIN (if any) before issuing a handout.
    gate = _pin_gate(request_id, item, pin)
    if gate is not None:
        return gate

    # Anti-hammer ceiling on the number of presigned handouts. Atomic: the same
    # op that counts the attempt rejects it once the ceiling is reached.
    try:
        _ddb.update_item(
            TableName=TABLE,
            Key={"id": {"S": request_id}},
            UpdateExpression="SET upload_attempts = if_not_exists(upload_attempts, :z) + :one",
            ConditionExpression="attribute_not_exists(upload_attempts) OR upload_attempts < :max",
            ExpressionAttributeValues={
                ":one": {"N": "1"},
                ":z": {"N": "0"},
                ":max": {"N": str(MAX_UPLOAD_ATTEMPTS)},
            },
        )
    except _ddb.exceptions.ConditionalCheckFailedException:
        return _resp(429, json.dumps({"error": "rate-limited"}), "application/json", _CORS)

    key = f"req/{request_id}"
    # A POST policy (not a bare PUT) so S3 enforces the size cap at upload time —
    # the gate never learns the size, so it couldn't sign a Content-Length itself.
    post = _s3.generate_presigned_post(
        Bucket=BUCKET,
        Key=key,
        Conditions=[["content-length-range", 0, MAX_UPLOAD]],
        ExpiresIn=900,
    )
    body = json.dumps({"url": post["url"], "fields": post["fields"], "key": key})
    return _resp(200, body, "application/json", _CORS)


def _done(request_id, pin, body_meta):
    """Finalize an upload the browser reports complete: confirm the object exists,
    record it, and spend one unit of the upload budget. Idempotent — a second
    call (retry, double-submit) is a no-op that still returns 200."""
    now = int(time.time())
    item = _ddb.get_item(TableName=TABLE, Key={"id": {"S": request_id}}).get("Item")
    if not item or item.get("kind", {}).get("S") != "request":
        return _resp(404, json.dumps({"error": "not found"}), "application/json", _CORS)

    # Same PIN gate as /up (the id is already MAC-gated).
    gate = _pin_gate(request_id, item, pin)
    if gate is not None:
        return gate

    key = f"req/{request_id}"
    try:
        head = _s3.head_object(Bucket=BUCKET, Key=key)
    except Exception:  # noqa: BLE001 - any head failure means "not confirmably there"
        return _resp(410, json.dumps({"error": "no upload found"}), "application/json", _CORS)
    size = int(head.get("ContentLength", 0))

    # Extract the encrypted filename blob from the JSON body (browser-produced).
    upload_meta = ""
    try:
        upload_meta = (json.loads(body_meta) or {}).get("upload_meta", "") if body_meta else ""
    except Exception:  # noqa: BLE001
        upload_meta = ""

    try:
        # Atomic + idempotent: record the object and spend one upload only if this
        # request has not already received one (attribute_not_exists(s3_key)).
        _ddb.update_item(
            TableName=TABLE,
            Key={"id": {"S": request_id}},
            UpdateExpression=(
                "SET s3_key = :k, size = :sz, upload_meta = :um, received_at = :now, "
                "uploads_remaining = uploads_remaining - :one"
            ),
            ConditionExpression="attribute_not_exists(s3_key)",
            ExpressionAttributeValues={
                ":k": {"S": key},
                ":sz": {"N": str(size)},
                ":um": {"S": upload_meta},
                ":now": {"N": str(now)},
                ":one": {"N": "1"},
            },
        )
    except _ddb.exceptions.ConditionalCheckFailedException:
        # Already received — a retry or a double-submit. Treat as success.
        return _resp(200, json.dumps({"ok": True}), "application/json", _CORS)

    return _resp(200, json.dumps({"ok": True}), "application/json", _CORS)