smolmail/smolmail.py

#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["noiseprotocol>=0.3.1", "pynacl>=1.5"]
# ///
"""Smol Mail reference client.

Implements SPEC.md version 1.1 from both ends: one master secret, sealed and
signed messages, server key pinning, trust on first use with rotation chains,
accept tokens, and the five operations.

Mail is stored sealed and opened on demand, so a stolen database is no more
readable than a stolen mailbox on a server.

    uv run smolmail.py keygen
    uv run smolmail.py trust example.org <server-key>
    uv run smolmail.py register alice@example.org
    uv run smolmail.py send bob@example.org --subject Hello
    uv run smolmail.py fetch
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import os
import re
import socket
import sqlite3
import struct
import sys
import time

import nacl.bindings as sodium
import nacl.exceptions
from nacl.signing import SigningKey, VerifyKey
from noise.connection import NoiseConnection

NOISE_PROTOCOL = b"Noise_NX_25519_ChaChaPoly_SHA256"
PROLOGUE = b"smolmail/1"
LABEL_AUTH, LABEL_SEAL = b"smolmail/1 auth", b"smolmail/1 seal"
LABEL_MSG, LABEL_ID, LABEL_ROTATE = b"smolmail/1 msg", b"smolmail/1 id", b"smolmail/1 rotate"
LABEL_IDENTITY, LABEL_ACCEPT = b"smolmail/1 identity", b"smolmail/1 accept"
LABEL_MAC, LABEL_REGISTER = b"smolmail/1 mac", b"smolmail/1 register"

OP_AUTH, OP_RESOLVE, OP_SEND, OP_FETCH, OP_DELETE, OP_REGISTER = range(6)
STATUS = {0: "ok", 1: "malformed", 2: "bad version", 3: "unknown user",
          4: "auth required", 5: "auth failed", 6: "quota exceeded", 7: "too large",
          8: "rate limited", 9: "not permitted", 10: "internal error"}

MAGIC, VERSION = b"SMOL", 1
KEY_LEN, ID_LEN, SIG_LEN, CERT_LEN, TOKEN_LEN = 32, 32, 64, 200, 32
PAYLOAD_HEADER = 45  # version 1 + sender 32 + time 8 + body_len 4
ENVELOPE_HEADER = 69  # magic 4 + version 1 + to 32 + epk 32
MAX_CHAIN, PAD_TO = 16, 1024
DEFAULT_PORT, MAX_FRAME, NOISE_PAYLOAD = 1961, 1 << 20, 65535 - 16
FRONTMATTER_MAX, FRONTMATTER_KEYS = 4096, 64
SEPARATORS = "._-"
MAX_SKEW = 86400  # §5.3, how far ahead of our clock a payload may be dated
TIER_MAIN, TIER_REQUESTS, FLAG_REQUESTS = 0, 1, 0x01


class SmolError(Exception):
    """Anything the user should see as a message rather than a traceback."""


# --- encoding ---------------------------------------------------------------


def b32(raw: bytes) -> str:
    """RFC 4648 base32, lowercase and unpadded (§3)."""
    return base64.b32encode(raw).decode("ascii").rstrip("=").lower()


def unb32(text: str) -> bytes:
    text = text.strip().upper()
    return base64.b32decode(text + "=" * (-len(text) % 8))


def fingerprint(identity: bytes) -> str:
    """First 20 characters of the base32 identity, in groups of four (§3)."""
    s = b32(identity)[:20]
    return " ".join(s[i : i + 4] for i in range(0, 20, 4))


class Reader:
    """Fail-closed reader; every parse raises rather than reading past the end."""

    def __init__(self, buf: bytes) -> None:
        self.buf, self.pos = buf, 0

    def take(self, n: int) -> bytes:
        if n < 0 or self.pos + n > len(self.buf):
            raise SmolError("truncated response from server")
        self.pos += n
        return self.buf[self.pos - n : self.pos]

    def u8(self) -> int:
        return self.take(1)[0]

    def u16(self) -> int:
        return struct.unpack(">H", self.take(2))[0]

    def u32(self) -> int:
        return struct.unpack(">I", self.take(4))[0]

    def i64(self) -> int:
        return struct.unpack(">q", self.take(8))[0]

    def left(self) -> int:
        return len(self.buf) - self.pos


ADDRESS = re.compile(r"^(?P<user>[a-z0-9._-]{1,63})@(?P<host>[^/:]+)(?::(?P<port>\d+))?$")


def valid_username(name: str) -> bool:
    """§3: alphanumeric at both ends, never two separators in a row."""
    if name[0] in SEPARATORS or name[-1] in SEPARATORS:
        return False
    return not any(a in SEPARATORS and b in SEPARATORS for a, b in zip(name, name[1:]))


class Address:
    """A short `user@host` address, or a self-certifying `smol://` one (§3)."""

    def __init__(self, user: str, host: str, port: int, identity: bytes | None = None) -> None:
        self.user, self.host, self.port, self.identity = user, host, port, identity

    @classmethod
    def parse(cls, text: str) -> "Address":
        text, identity = text.strip(), None
        if text.startswith("smol://"):
            rest = text[len("smol://") :]
            if "/" not in rest:
                raise SmolError(f"{text}: smol:// address carries no key")
            rest, key = rest.rsplit("/", 1)
            try:
                identity = unb32(key)
            except Exception as exc:
                raise SmolError(f"{text}: undecodable key: {exc}") from None
            if len(identity) != KEY_LEN:
                raise SmolError(f"{text}: key is {len(identity)} bytes, expected {KEY_LEN}")
            text = rest
        m = ADDRESS.match(text)
        if not m:
            raise SmolError(f"{text!r} is not a valid address")
        if not valid_username(m["user"]):
            raise SmolError(f"{m['user']!r} must begin and end with a letter or digit "
                            f"and may not contain two separators in a row")
        return cls(m["user"], m["host"], int(m["port"] or DEFAULT_PORT), identity)

    @property
    def short(self) -> str:
        return f"{self.user}@{self.host}" + ("" if self.port == DEFAULT_PORT else f":{self.port}")

    def uri(self, identity: bytes) -> str:
        return f"smol://{self.short}/{b32(identity)}"


# --- identity and message cryptography (§2, §5) -----------------------------


def hkdf_sha256(ikm: bytes, salt: bytes, info: bytes, length: int = 32) -> bytes:
    """RFC 5869, small enough to carry inline rather than add a dependency for."""
    prk = hmac.new(salt, ikm, hashlib.sha256).digest()
    out, block, counter = b"", b"", 1
    while len(out) < length:
        block = hmac.new(prk, block + info + bytes([counter]), hashlib.sha256).digest()
        out, counter = out + block, counter + 1
    return out[:length]


def identity_seed(master: bytes, index: int) -> bytes:
    """§2. The Ed25519 seed at a rotation index."""
    return hkdf_sha256(master, b"", LABEL_IDENTITY + struct.pack(">I", index))


def accept_key(master: bytes) -> bytes:
    """§2. Independent of the rotation index, so tokens survive a rotation."""
    return hkdf_sha256(master, b"", LABEL_ACCEPT)


def accept_mac(token: bytes, mid: bytes) -> bytes:
    """§5.8. What a sender attaches to SEND to reach the main tier."""
    return hmac.new(token, LABEL_MAC + mid, hashlib.sha256).digest()


def ed_to_x25519_pub(identity: bytes) -> bytes:
    """§2, via libsodium's own conversion as the spec prefers."""
    try:
        return sodium.crypto_sign_ed25519_pk_to_curve25519(identity)
    except (nacl.exceptions.CryptoError, RuntimeError, ValueError) as exc:
        raise SmolError(f"not a valid Ed25519 public key: {exc}") from None


class Identity:
    """An Ed25519 keypair plus the X25519 keypair derived from it (§2)."""

    def __init__(self, seed: bytes) -> None:
        if len(seed) != KEY_LEN:
            raise SmolError(f"identity seed must be {KEY_LEN} bytes, got {len(seed)}")
        self.seed = seed
        self._signing = SigningKey(seed)
        self.pk: bytes = bytes(self._signing.verify_key)
        self.x_priv: bytes = sodium.crypto_sign_ed25519_sk_to_curve25519(
            sodium.crypto_sign_seed_keypair(seed)[1])

    def sign(self, message: bytes) -> bytes:
        return self._signing.sign(message).signature


class Account:
    """The master secret of §2 and everything derived from it.

    `keys` holds every index up to the current one: sealing is to a specific key,
    so mail addressed to a superseded one is readable only with that key (§7).
    Rotation advances the index, so none of them has to be archived.
    """

    def __init__(self, master: bytes, index: int) -> None:
        if len(master) != KEY_LEN:
            raise SmolError(f"master secret must be {KEY_LEN} bytes, got {len(master)}")
        self.master, self.index = master, index
        self.keys = [Identity(identity_seed(master, n)) for n in range(index + 1)]
        self.me = self.keys[-1]

    def token_for(self, identity: bytes) -> bytes:
        """§5.8. The accept token this account issues to one correspondent."""
        return hmac.new(accept_key(self.master), identity, hashlib.sha256).digest()


def verify_sig(identity: bytes, signature: bytes, message: bytes) -> bool:
    try:
        VerifyKey(identity).verify(message, signature)
        return True
    except (nacl.exceptions.BadSignatureError, ValueError):
        return False


def agree(x_priv: bytes, x_pub: bytes) -> bytes:
    """X25519 with §2's checks; libsodium rejects low-order points itself."""
    try:
        shared = sodium.crypto_scalarmult(x_priv, x_pub)
    except (RuntimeError, nacl.exceptions.CryptoError) as exc:
        raise SmolError(f"rejected key agreement: {exc}") from None
    if shared == bytes(KEY_LEN):
        raise SmolError("rejected all-zero key agreement output")
    return shared


def message_id(envelope: bytes) -> bytes:
    """§5.4, derived from the envelope so no sender can choose it."""
    return hashlib.sha256(LABEL_ID + envelope).digest()


def seal(sender: Identity, recipient: bytes, body: bytes, when: int | None = None,
         pad: bool = True) -> bytes:
    """§5.2 and §5.3. Ephemeral-static, so the nonce is fixed and the sender's
    long-term key never takes part in key agreement."""
    when = int(time.time()) if when is None else when
    esk = os.urandom(KEY_LEN)
    epk = sodium.crypto_scalarmult_base(esk)
    key = hkdf_sha256(agree(esk, ed_to_x25519_pub(recipient)), epk + recipient, LABEL_SEAL)
    del esk  # best effort; CPython offers no way to wipe bytes

    header = bytes([VERSION]) + sender.pk + struct.pack(">qI", when, len(body))
    plaintext = header + body + sender.sign(LABEL_MSG + recipient + epk + header + body)
    if pad:
        plaintext += bytes(-len(plaintext) % PAD_TO)
    aad = MAGIC + bytes([VERSION]) + recipient + epk
    return aad + sodium.crypto_aead_chacha20poly1305_ietf_encrypt(
        plaintext, aad, bytes(12), key)


def unseal(identities: list[Identity], envelope: bytes) -> dict:
    """Inverse of seal(); raises unless signature and recipient both check out."""
    if len(envelope) < ENVELOPE_HEADER + 16:
        raise SmolError("envelope too short")
    if envelope[:4] != MAGIC:
        raise SmolError("not a Smol Mail envelope")
    if envelope[4] != VERSION:
        raise SmolError(f"unsupported envelope version {envelope[4]}")
    to, epk, sealed = envelope[5:37], envelope[37:69], envelope[69:]

    me = next((i for i in identities if i.pk == to), None)
    if me is None:
        raise SmolError(f"addressed to {b32(to)[:16]}…, not one of our keys")
    key = hkdf_sha256(agree(me.x_priv, epk), epk + to, LABEL_SEAL)
    try:
        plaintext = sodium.crypto_aead_chacha20poly1305_ietf_decrypt(
            sealed, envelope[:ENVELOPE_HEADER], bytes(12), key)
    except nacl.exceptions.CryptoError:
        raise SmolError("decryption failed: wrong key or corrupt envelope") from None

    r = Reader(plaintext)
    if r.u8() != VERSION:
        raise SmolError("unsupported payload version")
    sender, when, body_len = r.take(KEY_LEN), r.i64(), r.u32()
    if body_len > r.left():
        raise SmolError("payload body length exceeds the payload")
    body, signature = r.take(body_len), r.take(SIG_LEN)  # trailing bytes are padding
    if not verify_sig(sender, signature,
                      LABEL_MSG + to + epk + plaintext[:PAYLOAD_HEADER] + body):
        raise SmolError("signature does not verify")
    if when > int(time.time()) + MAX_SKEW:
        raise SmolError("payload is dated in the future")
    return {"sender": sender, "time": when, "body": body, "id": message_id(envelope)}


# --- body frontmatter (§5.5) ------------------------------------------------

FM_KEY = re.compile(r"^[A-Za-z0-9-]{1,64}$")


def parse_frontmatter(body: str) -> tuple[dict[str, str], str]:
    """§5.5. A flat `Key: value` block, deliberately not YAML.

    Any malformed line invalidates the whole block, which is then returned as
    ordinary body text: frontmatter fails closed toward display. Keys are
    compared case-insensitively, so they are kept lowercased.
    """
    if not body.startswith("---\n"):
        return {}, body
    lines = body.split("\n")
    try:
        close = lines.index("---", 1)
    except ValueError:
        return {}, body
    block, rest = lines[1:close], "\n".join(lines[close + 1 :])
    if len(block) > FRONTMATTER_KEYS:
        return {}, body
    if sum(len(line.encode()) + 1 for line in block) > FRONTMATTER_MAX:
        return {}, body
    fields: dict[str, str] = {}
    for line in block:
        key, sep, value = line.partition(":")
        if not sep or not FM_KEY.match(key):
            return {}, body
        fields.setdefault(key.lower(), value.strip())  # first occurrence wins
    return fields, rest


def build_frontmatter(fields: list[tuple[str, str]], body: str) -> str:
    """Emit a block only when needed, including to escape a body that genuinely
    begins with `---` (§5.5)."""
    if not fields and not body.startswith("---\n"):
        return body
    return "---\n" + "".join(f"{k}: {v}\n" for k, v in fields) + "---\n" + body


# --- local state ------------------------------------------------------------

SCHEMA = """
-- One row, always present, so every update is a plain UPDATE.
CREATE TABLE IF NOT EXISTS state (
    id INTEGER PRIMARY KEY CHECK (id = 1),
    username TEXT, host TEXT, port INTEGER,
    rotations  INTEGER NOT NULL DEFAULT 0,   -- §2 rotation index
    after_time INTEGER NOT NULL DEFAULT 0,   -- §6.1 FETCH cursor
    after_id   BLOB    NOT NULL DEFAULT x'',
    sync_ok    INTEGER NOT NULL DEFAULT 1);  -- may we replace the server's set?
INSERT OR IGNORE INTO state (id) VALUES (1);
CREATE TABLE IF NOT EXISTS servers (
    host TEXT PRIMARY KEY, static BLOB NOT NULL, pinned_at INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS contacts (
    address TEXT PRIMARY KEY, identity BLOB NOT NULL,
    verified INTEGER NOT NULL,          -- 1 when the key came from a smol:// address
    seen_at INTEGER NOT NULL);
-- Correspondents admitted to this mailbox's main tier (§5.8). The identity is
-- frozen at acceptance because the token is derived from it: a contact's later
-- rotation must not change the token they already hold.
CREATE TABLE IF NOT EXISTS accepted (
    address TEXT PRIMARY KEY, identity BLOB NOT NULL,
    active INTEGER NOT NULL, added_at INTEGER NOT NULL);
-- Accept tokens received from correspondents, filed under the address that
-- issued them: an address outlives the keys behind it, so a token keeps working
-- across the issuer's rotations (§5.8).
CREATE TABLE IF NOT EXISTS tokens (
    address TEXT PRIMARY KEY, token BLOB NOT NULL, seen_at INTEGER NOT NULL);
-- Every id ever fetched, so an envelope resent after we deleted it locally is
-- not stored again (§10).
CREATE TABLE IF NOT EXISTS seen (id BLOB PRIMARY KEY, at INTEGER NOT NULL);
-- Envelopes are stored sealed and opened on demand; no plaintext at rest.
CREATE TABLE IF NOT EXISTS inbox (
    id BLOB PRIMARY KEY, envelope BLOB NOT NULL,
    received_at INTEGER NOT NULL, tier INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS sent (
    id BLOB PRIMARY KEY, recipient TEXT NOT NULL,
    envelope BLOB NOT NULL, sent_at INTEGER NOT NULL);
"""


class Store:
    def __init__(self, path: str) -> None:
        self.db = sqlite3.connect(path)
        self.db.executescript(SCHEMA)
        self.db.commit()

    def all(self, sql: str, *args) -> list[tuple]:
        return self.db.execute(sql, args).fetchall()

    def one(self, sql: str, *args) -> tuple | None:
        return self.db.execute(sql, args).fetchone()

    def run(self, sql: str, *args) -> sqlite3.Cursor:
        with self.db:
            return self.db.execute(sql, args)

    def account(self) -> Address | None:
        row = self.one("SELECT username, host, port FROM state WHERE id = 1")
        return Address(row[0], row[1], row[2]) if row and row[0] else None

    def set_account(self, addr: Address) -> None:
        self.run("UPDATE state SET username = ?, host = ?, port = ? WHERE id = 1",
                 addr.user, addr.host, addr.port)

    def rotations(self) -> int:
        return self.one("SELECT rotations FROM state WHERE id = 1")[0]

    def cursor(self) -> tuple[int, bytes]:
        after_time, after_id = self.one("SELECT after_time, after_id FROM state WHERE id = 1")
        return after_time, after_id if len(after_id) == ID_LEN else bytes(ID_LEN)

    def set_cursor(self, after_time: int, after_id: bytes) -> None:
        self.run("UPDATE state SET after_time = ?, after_id = ? WHERE id = 1",
                 after_time, after_id)

    def contact(self, address: str) -> tuple[bytes, bool] | None:
        row = self.one("SELECT identity, verified FROM contacts WHERE address = ?", address)
        return (row[0], bool(row[1])) if row else None

    def save_contact(self, address: str, identity: bytes, verified: bool) -> None:
        self.run("INSERT INTO contacts (address, identity, verified, seen_at) VALUES (?,?,?,?) "
                 "ON CONFLICT (address) DO UPDATE SET identity = ?2, verified = ?3, seen_at = ?4",
                 address, identity, int(verified), int(time.time()))

    def pin_server(self, host: str, static: bytes) -> None:
        self.run("INSERT INTO servers (host, static, pinned_at) VALUES (?,?,?) "
                 "ON CONFLICT (host) DO UPDATE SET static = ?2, pinned_at = ?3",
                 host, static, int(time.time()))

    def token_set(self, account: "Account") -> tuple[int, list[bytes]]:
        """§4. The accept tokens to push with AUTH, and whether to push at all.

        A client that cannot vouch for its own set — one restored from the master
        alone — must not replace the server's with an incomplete one.
        """
        if not self.one("SELECT sync_ok FROM state WHERE id = 1")[0]:
            return 0, []
        rows = self.all("SELECT identity FROM accepted WHERE active = 1 ORDER BY added_at")
        return 1, [account.token_for(row[0]) for row in rows]

    def address_of(self, sender: bytes, reply_to: str | None) -> str | None:
        """The address we know a signer by: a contact, or the Reply-To it signed
        for itself. Naming a mailbox is not trusting a key, so nothing is
        pinned here (§5.7, §8)."""
        for address, identity in self.all("SELECT address, identity FROM contacts"):
            if hmac.compare_digest(identity, sender):
                return address
        if reply_to:
            try:
                claimed = Address.parse(reply_to)
            except SmolError:
                return None
            if claimed.identity and hmac.compare_digest(claimed.identity, sender):
                return claimed.short
        return None

    def learn_token(self, opened: dict) -> None:
        """§5.8. An Accept field is bound to the signer of the message that
        carried it, which unseal() has already verified."""
        fields, _ = parse_frontmatter(opened["body"].decode("utf-8", "replace"))
        raw = fields.get("accept")
        if not raw:
            return
        try:
            token = unb32(raw)
        except Exception:
            return
        if len(token) != TOKEN_LEN:
            return
        address = self.address_of(opened["sender"], fields.get("reply-to"))
        if address is None:
            return  # no address to send to, so no use for a token
        self.run("INSERT INTO tokens (address, token, seen_at) VALUES (?,?,?) "
                 "ON CONFLICT (address) DO UPDATE SET token = ?2, seen_at = ?3",
                 address, token, int(time.time()))

    def mail(self, box: str) -> list[tuple]:
        if box == "sent":
            return self.all("SELECT id, envelope, sent_at, recipient FROM sent "
                            "ORDER BY sent_at, id")
        if box == "all":
            return self.all("SELECT id, envelope, received_at, NULL FROM inbox "
                            "ORDER BY received_at, id")
        tier = TIER_REQUESTS if box == "requests" else TIER_MAIN
        return self.all("SELECT id, envelope, received_at, NULL FROM inbox "
                        "WHERE tier = ? ORDER BY received_at, id", tier)


# --- transport (§4) ---------------------------------------------------------


class Conn:
    """A Noise_NX initiator session with pinning layered on top (§4)."""

    def __init__(self, host: str, port: int, pinned: bytes | None) -> None:
        try:
            self.sock = socket.create_connection((host, port), timeout=30)
        except OSError as exc:
            raise SmolError(f"cannot reach {host}:{port}: {exc}") from None
        noise = NoiseConnection.from_name(NOISE_PROTOCOL)
        noise.set_as_initiator()
        noise.set_prologue(PROLOGUE)
        noise.start_handshake()
        first = noise.write_message()
        self.sock.sendall(struct.pack(">H", len(first)) + first)
        # The library drops the peer's static key once the handshake ends, so
        # keep the handshake state to read it afterwards.
        state = noise.noise_protocol.handshake_state
        noise.read_message(self._recv(struct.unpack(">H", self._recv(2))[0]))
        if not noise.handshake_finished:
            raise SmolError("handshake did not complete")

        self.server_static = bytes(state.rs.public_bytes)
        if pinned is not None and not hmac.compare_digest(pinned, self.server_static):
            self.close()
            raise SmolError(f"{host} presented a different key than the one pinned\n"
                            f"  pinned:    {b32(pinned)}\n"
                            f"  presented: {b32(self.server_static)}")
        self.pinned = pinned is not None
        self.noise, self.buf = noise, bytearray()
        self.handshake_hash: bytes = noise.get_handshake_hash()

    def _recv(self, n: int) -> bytes:
        out = b""
        while len(out) < n:
            chunk = self.sock.recv(n - len(out))
            if not chunk:
                raise SmolError("server closed the connection")
            out += chunk
        return out

    def call(self, op: int, body: bytes = b"") -> tuple[int, bytes]:
        frame = struct.pack(">I", 1 + len(body)) + bytes([op]) + body
        if len(frame) > MAX_FRAME + 4:
            raise SmolError("request exceeds the maximum frame size")
        for off in range(0, len(frame), NOISE_PAYLOAD):
            packet = self.noise.encrypt(frame[off : off + NOISE_PAYLOAD])
            self.sock.sendall(struct.pack(">H", len(packet)) + packet)
        while len(self.buf) < 5:
            self.buf += self.noise.decrypt(self._recv(struct.unpack(">H", self._recv(2))[0]))
        (length,) = struct.unpack(">I", self.buf[:4])
        if not 1 <= length <= MAX_FRAME:
            raise SmolError(f"server sent a frame of length {length}")
        while len(self.buf) < 4 + length:
            self.buf += self.noise.decrypt(self._recv(struct.unpack(">H", self._recv(2))[0]))
        payload = bytes(self.buf[4 : 4 + length])
        del self.buf[: 4 + length]
        return payload[1], payload[2:]  # status, body (§6.1)

    def close(self) -> None:
        try:
            self.sock.close()
        except OSError:
            pass

    def __enter__(self) -> "Conn":
        return self

    def __exit__(self, *exc) -> None:
        self.close()


def connect(store: Store, addr: Address, require_pin: bool) -> Conn:
    """Open a session, enforcing §4's rule about unpinned servers."""
    row = store.one("SELECT static FROM servers WHERE host = ?", addr.host)
    pinned = row[0] if row else None
    if pinned is None and require_pin:
        raise SmolError(f"no pinned key for {addr.host}.\nObtain it from the operator "
                        f"through a trusted channel, then:\n"
                        f"  smolmail.py trust {addr.host} <key>")
    conn = Conn(addr.host, addr.port, pinned)
    if pinned is None:
        warn(f"{addr.host} is not pinned; its key is {b32(conn.server_static)}")
    return conn


def expect_ok(status: int, what: str) -> None:
    if status != 0:
        raise SmolError(f"{what} failed: {STATUS.get(status, status)} ({status})")


# --- operations -------------------------------------------------------------


def do_resolve(conn: Conn, user: str) -> tuple[bytes, list[bytes]]:
    """RESOLVE, returning the current key and its rotation chain (§6.1)."""
    name = user.encode()
    status, body = conn.call(OP_RESOLVE, bytes([len(name)]) + name)
    expect_ok(status, f"resolving {user}")
    r = Reader(body)
    identity = r.take(KEY_LEN)
    return identity, [r.take(CERT_LEN) for _ in range(r.u8())]


def walk_chain(username: str, pinned: bytes, current: bytes, chain: list[bytes]) -> bool:
    """§7. Accept a key change only when a signed chain leads from the key we
    hold to the one the server now returns. Both keys must sign each link."""
    if hmac.compare_digest(pinned, current):
        return True
    if not chain or len(chain) > MAX_CHAIN:
        return False
    key, started = pinned, False
    for cert in chain:
        old, new, when = cert[:32], cert[32:64], cert[64:72]
        sig_old, sig_new = cert[72:136], cert[136:200]
        if not started:
            if old != key:
                continue  # a link predating the key we hold
            started = True
        elif old != key:
            return False  # the chain is not continuous
        signed = LABEL_ROTATE + username.encode() + old + new + when
        if not verify_sig(old, sig_old, signed) or not verify_sig(new, sig_new, signed):
            return False
        key = new
    return started and hmac.compare_digest(key, current)


def trust_key(store: Store, conn: Conn, addr: Address) -> bytes:
    """Resolve a contact and apply §8's trust rules."""
    identity, chain = do_resolve(conn, addr.user)
    known = store.contact(addr.short)
    if known is None:
        store.save_contact(addr.short, identity, verified=False)
        qualifier = "" if conn.pinned else ", UNVERIFIED server"
        print(f"{addr.short}  {b32(identity)}\n  pinned (trust on first use{qualifier})")
        return identity
    old, verified = known
    if hmac.compare_digest(old, identity):
        return identity
    if walk_chain(addr.user, old, identity, chain):
        store.save_contact(addr.short, identity, verified)
        warn(f"{addr.short} rotated its key; a signed chain confirms it")
        print(f"  now {b32(identity)}")
        return identity
    raise SmolError(f"{addr.short} presents a different key with no valid rotation chain.\n"
                    f"  known:   {b32(old)}\n  offered: {b32(identity)}\n"
                    f"Verify out of band, then: smolmail.py import {addr.uri(identity)}")


def register_signed(server_static: bytes, username: str, identity: bytes) -> bytes:
    """§6.1 proof of possession, bound to the server that will store the binding."""
    return LABEL_REGISTER + server_static + username.encode() + identity


def authenticate(conn: Conn, addr: Address, account: Account, store: Store) -> int:
    """§4 session authentication: sign the handshake hash and push the token set."""
    name = addr.user.encode()
    sync, tokens = store.token_set(account)
    body = (bytes([len(name)]) + name + account.me.pk
            + account.me.sign(LABEL_AUTH + conn.handshake_hash)
            + bytes([sync]) + struct.pack(">H", len(tokens)) + b"".join(tokens))
    status, payload = conn.call(OP_AUTH, body)
    expect_ok(status, "authentication")
    return struct.unpack(">H", payload[:2])[0] if len(payload) >= 2 else 0


def push_tokens(store: Store, account: Account) -> None:
    """§5.8. An accept or a block only takes effect once the server holds the
    changed set, so it is pushed now rather than at the next fetch."""
    addr = store.account()
    if addr is None:
        warn("not registered; the set will be pushed with your first fetch")
        return
    with connect(store, addr, require_pin=True) as conn:
        held = authenticate(conn, addr, account, store)
    print(f"server now holds {held} accept token(s)")


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


def warn(message: str) -> None:
    print(f"warning: {message}", file=sys.stderr)


def load_master(args: argparse.Namespace) -> bytes:
    try:
        with open(args.key, "rb") as fh:
            return fh.read()
    except FileNotFoundError:
        raise SmolError(f"no identity at {args.key}; run: smolmail.py keygen") from None


def load_account(args: argparse.Namespace, store: Store) -> Account:
    """§2. The master from disk plus the rotation index from local state."""
    return Account(load_master(args), store.rotations())


def write_secret(path: str, secret: bytes) -> None:
    """0600 before any bytes land, so the secret is never briefly world-readable."""
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "wb") as fh:
        fh.write(secret)


def describe(store: Store, envelope: bytes, identities: list[Identity]) -> dict:
    """Open an envelope and split its body into frontmatter and text."""
    opened = unseal(identities, envelope)
    fields, text = parse_frontmatter(opened["body"].decode("utf-8", "replace"))
    sender = opened["sender"]
    known = next((a for a, i, _ in store.all(
        "SELECT address, identity, verified FROM contacts") if hmac.compare_digest(i, sender)), None)
    opened |= {"fields": fields, "text": text, "subject": fields.get("subject", ""),
               "from": known or f"<{b32(sender)[:20]}…>"}
    return opened


def target_contact(store: Store, text: str) -> tuple[str, bytes]:
    """An address plus the key we hold for it, for commands naming a contact."""
    addr = Address.parse(text)
    if addr.identity is not None:
        return addr.short, addr.identity
    known = store.contact(addr.short)
    if known is None:
        raise SmolError(f"no key for {addr.short}; run `resolve` or `import` first")
    return addr.short, known[0]


# --- commands ---------------------------------------------------------------


def cmd_keygen(args: argparse.Namespace, store: Store) -> int:
    if os.path.exists(args.key) and not args.force:
        raise SmolError(f"{args.key} exists; refusing to overwrite (use --force)")
    master = os.urandom(KEY_LEN)
    write_secret(args.key, master)
    account = Account(master, 0)
    print(f"master:      {args.key}  (back this up; it is the only secret)")
    print(f"public key:  {b32(account.me.pk)}\nfingerprint: {fingerprint(account.me.pk)}")
    return 0


def cmd_whoami(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    print(f"public key:  {b32(account.me.pk)}\nfingerprint: {fingerprint(account.me.pk)}")
    addr = store.account()
    print(f"address:     {addr.short}\nuri:         {addr.uri(account.me.pk)}" if addr
          else "address:     (not registered)")
    if account.index:
        print(f"rotations:   {account.index} (earlier keys derived on demand)")
    live = store.one("SELECT COUNT(*) FROM accepted WHERE active = 1")[0]
    sync = store.one("SELECT sync_ok FROM state WHERE id = 1")[0]
    print(f"accepted:    {live} correspondent(s)"
          + ("" if sync else ", not pushed to the server until rebuilt"))
    return 0


def cmd_trust(args: argparse.Namespace, store: Store) -> int:
    try:
        key = unb32(args.key_b32)
    except Exception as exc:
        raise SmolError(f"undecodable key: {exc}") from None
    if len(key) != KEY_LEN:
        raise SmolError(f"server key is {len(key)} bytes, expected {KEY_LEN}")
    host = args.host.split(":")[0]
    row = store.one("SELECT static FROM servers WHERE host = ?", host)
    if row and row[0] != key and not args.force:
        raise SmolError(f"{host} is already pinned to {b32(row[0])}; use --force to replace")
    store.pin_server(host, key)
    print(f"pinned {host}  {b32(key)}")
    return 0


def cmd_register(args: argparse.Namespace, store: Store) -> int:
    account, addr = load_account(args, store), Address.parse(args.address)
    me = account.me
    name, token = addr.user.encode(), (args.token or "").encode()
    with connect(store, addr, require_pin=True) as conn:
        body = (bytes([len(name)]) + name + me.pk
                + me.sign(register_signed(conn.server_static, addr.user, me.pk))
                + bytes([len(token)]) + token + b"\0")
        status, _ = conn.call(OP_REGISTER, body)
    expect_ok(status, f"registering {addr.short}")
    store.set_account(addr)
    print(f"registered {addr.short}\nshare: {addr.uri(me.pk)}")
    return 0


def cmd_resolve(args: argparse.Namespace, store: Store) -> int:
    addr = Address.parse(args.address)
    if addr.identity is not None:
        raise SmolError("that address already carries a key; use `import` instead")
    with connect(store, addr, require_pin=False) as conn:
        trust_key(store, conn, addr)
    return 0


def cmd_import(args: argparse.Namespace, store: Store) -> int:
    addr = Address.parse(args.uri)
    if addr.identity is None:
        raise SmolError("import needs a smol:// address carrying a key")
    store.save_contact(addr.short, addr.identity, verified=True)
    print(f"imported {addr.short}  {b32(addr.identity)}  (verified)")
    print(f"fingerprint: {fingerprint(addr.identity)}")
    return 0


def cmd_contacts(args: argparse.Namespace, store: Store) -> int:
    rows = store.all("SELECT c.address, c.identity, c.verified, a.active FROM contacts c "
                     "LEFT JOIN accepted a ON a.address = c.address ORDER BY c.address")
    if not rows:
        print("no contacts")
        return 0
    width = max(len(a) for a, _, _, _ in rows)
    for address, identity, verified, active in rows:
        state = "accepted" if active else ("blocked" if active == 0 else "")
        print(f"{address:<{width}}  {b32(identity)}  "
              f"{'verified' if verified else 'tofu':<8}  {state}".rstrip())
    return 0


def cmd_accept(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    address, identity = target_contact(store, args.address)
    if not store.one("SELECT sync_ok FROM state WHERE id = 1")[0]:
        warn("this client's accepted set was not restored; from now on it replaces "
             "the server's, so re-accept everyone you still correspond with")
    # ON CONFLICT leaves `identity` alone: the token stays the one the
    # correspondent already holds, even after they rotate (§5.8).
    store.run("INSERT INTO accepted (address, identity, active, added_at) VALUES (?,?,1,?) "
              "ON CONFLICT (address) DO UPDATE SET active = 1", address, identity,
              int(time.time()))
    store.run("UPDATE state SET sync_ok = 1 WHERE id = 1")
    print(f"accepted {address}; its token travels in your next message to them")
    push_tokens(store, account)
    return 0


def cmd_block(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    address, _ = target_contact(store, args.address)
    if not store.run("UPDATE accepted SET active = 0 WHERE address = ?", address).rowcount:
        raise SmolError(f"{address} was never accepted")
    print(f"blocked {address}; their mail lands in requests from their next message on")
    push_tokens(store, account)
    return 0


def cmd_send(args: argparse.Namespace, store: Store) -> int:
    account, addr = load_account(args, store), Address.parse(args.address)
    me = account.me
    if args.body is not None:
        text = args.body
    elif sys.stdin.isatty():
        raise SmolError("no message body; pass --body or pipe it on stdin")
    else:
        text = sys.stdin.read()

    # Prefer a key we already trust; fall back to RESOLVE with trust on first use.
    if addr.identity is not None:
        recipient = addr.identity
        store.save_contact(addr.short, recipient, verified=True)
    elif (known := store.contact(addr.short)) is not None:
        recipient = known[0]
    else:
        with connect(store, addr, require_pin=False) as conn:
            recipient = trust_key(store, conn, addr)

    fields = [(k, v) for k, v in (("Subject", args.subject), ("In-Reply-To", args.reply_to)) if v]
    for raw in args.header or []:
        key, sep, value = raw.partition(":")
        if not sep or not FM_KEY.match(key.strip()):
            raise SmolError(f"{raw!r} is not a valid `Key: value` header")
        fields.append((key.strip(), value.strip()))
    # §5.7: a signed reply address lets a first-time recipient answer us.
    if (account_addr := store.account()) is not None and not args.anonymous:
        fields.append(("Reply-To", account_addr.uri(me.pk)))
    # §5.8: hand an accepted correspondent the token for our own mailbox.
    if (row := store.one("SELECT identity FROM accepted WHERE address = ? AND active = 1",
                         addr.short)) is not None:
        fields.append(("Accept", b32(account.token_for(row[0]))))

    body = build_frontmatter(fields, text).encode()
    envelope = seal(me, recipient, body, pad=not args.no_pad)
    mid = message_id(envelope)
    # §5.8: our token for their mailbox, if they have given us one.
    held = store.one("SELECT token FROM tokens WHERE address = ?", addr.short)
    mac = accept_mac(held[0], mid) if held else b""
    with connect(store, addr, require_pin=False) as conn:
        status, payload = conn.call(OP_SEND, bytes([len(mac)]) + mac + envelope)
    expect_ok(status, f"sending to {addr.short}")
    if payload and payload != mid:
        warn("server returned an id we did not derive; it is not authoritative")
    # §5.6: the ephemeral is gone, so keep a copy sealed to ourselves.
    store.run("INSERT OR IGNORE INTO sent (id, recipient, envelope, sent_at) VALUES (?,?,?,?)",
              mid, addr.short, seal(me, me.pk, body, pad=not args.no_pad), int(time.time()))
    print(f"sent {mid.hex()[:16]} to {addr.short} ({len(envelope)} bytes"
          + (", accepted" if mac else "") + ")")
    return 0


def cmd_fetch(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    addr = store.account()
    if addr is None:
        raise SmolError("not registered; run: smolmail.py register <user@host>")
    if args.reset:
        store.set_cursor(0, bytes(ID_LEN))
    # Acknowledging deletes what it takes, so it always pages from the start and
    # leaves the stored cursor at zero; --keep instead remembers its place.
    after_time, after_id = store.cursor() if args.keep else (0, bytes(ID_LEN))
    stored = rejected = total = 0

    with connect(store, addr, require_pin=True) as conn:
        authenticate(conn, addr, account, store)
        while True:
            status, body = conn.call(OP_FETCH, struct.pack(">q", after_time) + after_id)
            expect_ok(status, "fetching")
            r = Reader(body)
            count = r.u16()
            if count == 0:
                break
            acked = []
            for _ in range(count):
                total += 1
                mid, received_at, flags = r.take(ID_LEN), r.i64(), r.u8()
                envelope = r.take(r.u32())
                after_time, after_id = received_at, mid
                try:
                    if message_id(envelope) != mid:
                        raise SmolError("id does not match the envelope")
                    opened = unseal(account.keys, envelope)
                except SmolError as exc:
                    # Left on the server rather than destroyed, so a client-side
                    # bug cannot lose mail.
                    warn(f"{mid.hex()[:16]}: {exc}; left on server")
                    rejected += 1
                    continue
                # A message we have already had once is not stored again, even
                # if we deleted it locally in the meantime (§10).
                if store.one("SELECT 1 FROM seen WHERE id = ?", mid) is None:
                    tier = TIER_REQUESTS if flags & FLAG_REQUESTS else TIER_MAIN
                    store.run("INSERT OR IGNORE INTO inbox "
                              "(id, envelope, received_at, tier) VALUES (?,?,?,?)",
                              mid, envelope, received_at, tier)
                    store.run("INSERT OR IGNORE INTO seen (id, at) VALUES (?,?)",
                              mid, received_at)
                    store.learn_token(opened)
                    stored += 1
                acked.append(mid)
            if args.keep:
                store.set_cursor(after_time, after_id)
            elif acked:
                status, _ = conn.call(OP_DELETE,
                                      struct.pack(">H", len(acked)) + b"".join(acked))
                expect_ok(status, "acknowledging")

    if not args.keep:
        store.set_cursor(0, bytes(ID_LEN))
    print(f"{total} message(s): {stored} new, {rejected} rejected"
          + (" (left on the server)" if args.keep and total else ""))
    return 0


def cmd_list(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    box = "sent" if args.sent else "requests" if args.requests else "inbox"
    rows = store.mail(box)
    if not rows:
        print(f"{box} is empty")
        return 0
    for mid, envelope, when, recipient in rows:
        try:
            opened = describe(store, envelope, account.keys)
            who = recipient if box == "sent" else opened["from"]
            subject = opened["subject"] or "(no subject)"
        except SmolError as exc:
            who, subject = "?", f"<unreadable: {exc}>"
        stamp = time.strftime("%Y-%m-%d %H:%M", time.localtime(when))
        print(f"{mid.hex()[:8]}  {stamp}  {who:<28.28}  {subject}")
    return 0


def cmd_read(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    box = "sent" if args.sent else "all"
    matches = [r for r in store.mail(box) if r[0].hex().startswith(args.id.lower())]
    if not matches:
        raise SmolError(f"no {'sent ' if args.sent else ''}message matching {args.id!r}")
    if len(matches) > 1:
        raise SmolError(f"{args.id!r} matches {len(matches)} messages; be more specific")
    mid, envelope, _, recipient = matches[0]
    opened = describe(store, envelope, account.keys)
    def field(key: str, value: str) -> None:
        print(f"{key + ':':<12} {value}")

    field("id", mid.hex())
    if box == "sent":
        field("to", recipient)
    field("from", opened["from"])
    field("key", b32(opened["sender"]))
    field("date", time.strftime("%Y-%m-%d %H:%M:%S %z", time.localtime(opened["time"])))
    for key, value in opened["fields"].items():
        if key != "accept":  # machinery, not content (§5.8)
            field(key, value)
    print("signature verified\n")
    print(opened["text"], end="" if opened["text"].endswith("\n") else "\n")
    return 0


def cmd_rotate(args: argparse.Namespace, store: Store) -> int:
    account = load_account(args, store)
    addr = store.account()
    if addr is None:
        raise SmolError("not registered; run: smolmail.py register <user@host>")
    if account.index >= MAX_CHAIN:
        raise SmolError(f"the rotation chain is full at {MAX_CHAIN} links")
    old = account.me
    new = Identity(identity_seed(account.master, account.index + 1))
    when = struct.pack(">q", int(time.time()))
    # §7: both keys sign, so the old key alone cannot hand the username to a
    # key nobody controls.
    signed = LABEL_ROTATE + addr.user.encode() + old.pk + new.pk + when
    cert = old.pk + new.pk + when + old.sign(signed) + new.sign(signed)

    name = addr.user.encode()
    with connect(store, addr, require_pin=True) as conn:
        body = (bytes([len(name)]) + name + new.pk
                + new.sign(register_signed(conn.server_static, addr.user, new.pk))
                + b"\0" + bytes([len(cert)]) + cert)
        status, _ = conn.call(OP_REGISTER, body)
    expect_ok(status, "rotating")

    # The master is untouched; only the index moves, and the superseded key stays
    # derivable from it (§2).
    store.run("UPDATE state SET rotations = ? WHERE id = 1", account.index + 1)
    print(f"rotated {addr.short}\nnew key:     {b32(new.pk)}")
    print(f"fingerprint: {fingerprint(new.pk)}\nshare: {addr.uri(new.pk)}")
    print("\nTell your contacts; they will accept the change from the signed chain.")
    return 0


def cmd_restore(args: argparse.Namespace, store: Store) -> int:
    """§2. Recover the rotation index, and with it every superseded key, from
    the master alone."""
    master = load_master(args)
    addr = Address.parse(args.address)
    with connect(store, addr, require_pin=False) as conn:
        identity, _ = do_resolve(conn, addr.user)
    for index in range(MAX_CHAIN + 1):
        if Identity(identity_seed(master, index)).pk == identity:
            break
    else:
        raise SmolError(f"the key bound to {addr.short} is not derived from this master "
                        f"within {MAX_CHAIN} rotations")
    store.set_account(addr)
    store.set_cursor(0, bytes(ID_LEN))
    # The accepted set is gone, and an empty one must not replace the server's (§4).
    store.run("UPDATE state SET rotations = ?, sync_ok = 0 WHERE id = 1", index)
    print(f"restored {addr.short} at rotation {index}\npublic key:  {b32(identity)}")
    print("Your accepted correspondents are still live on the server; `accept` them "
          "again only when you are ready to replace that set.")
    return 0


# --- CLI --------------------------------------------------------------------

COMMANDS = [
    ("keygen", cmd_keygen, "create a master secret",
     [(("--force",), {"action": "store_true"})]),
    ("whoami", cmd_whoami, "show this identity", []),
    ("trust", cmd_trust, "pin a server's static key",
     [(("host",), {}), (("key_b32",), {"metavar": "key"}),
      (("--force",), {"action": "store_true"})]),
    ("register", cmd_register, "bind this identity to a username",
     [(("address",), {}), (("--token",), {"help": "invite token, if required"})]),
    ("restore", cmd_restore, "recover local state from the master",
     [(("address",), {})]),
    ("resolve", cmd_resolve, "look up and pin a contact's key", [(("address",), {})]),
    ("import", cmd_import, "add a contact from a smol:// address", [(("uri",), {})]),
    ("contacts", cmd_contacts, "list known keys", []),
    ("accept", cmd_accept, "admit a contact to the main tier", [(("address",), {})]),
    ("block", cmd_block, "withdraw a contact's accept token", [(("address",), {})]),
    ("send", cmd_send, "seal and deliver a message",
     [(("address",), {}), (("--subject",), {}),
      (("--reply-to",), {"metavar": "ID"}),
      (("--body",), {"help": "message text; read from stdin when omitted"}),
      (("--header",), {"action": "append", "metavar": "KEY:VALUE",
                       "help": "extra frontmatter field; repeatable"}),
      (("--anonymous",), {"action": "store_true",
                          "help": "omit the Reply-To field carrying this address (SPEC.md §5.7)"}),
      (("--no-pad",), {"action": "store_true", "help": "do not pad to 1 KiB"})]),
    ("fetch", cmd_fetch, "retrieve, verify and acknowledge mail",
     [(("--keep",), {"action": "store_true",
                     "help": "do not delete from the server; remember the cursor instead"}),
      (("--reset",), {"action": "store_true", "help": "forget the cursor and page again"})]),
    ("list", cmd_list, "list stored mail",
     [(("--sent",), {"action": "store_true"}),
      (("--requests",), {"action": "store_true",
                         "help": "mail that arrived without an accept token"})]),
    ("read", cmd_read, "show one message",
     [(("id",), {}), (("--sent",), {"action": "store_true"})]),
    ("rotate", cmd_rotate, "advance to the next identity, signing the change", []),
]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--key", default="identity.key", help="master secret file")
    parser.add_argument("--db", default="smolmail.db", help="local state and mail")
    sub = parser.add_subparsers(dest="command", required=True)
    for name, func, help_text, arguments in COMMANDS:
        child = sub.add_parser(name, help=help_text)
        child.set_defaults(func=func)
        for flags, options in arguments:
            child.add_argument(*flags, **options)

    args = parser.parse_args(argv)
    try:
        return args.func(args, Store(args.db))
    except SmolError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        return 130


if __name__ == "__main__":
    sys.exit(main())