#!/usr/bin/env python3
"""
kikemail.py - create addresses and read encrypted mail without the website.

No dependencies beyond Python 3.8+. All crypto is implemented here in pure
Python so the entire trust path is auditable in a single file:

    X25519 (RFC 7748) + HSalsa20/XSalsa20-Poly1305 (NaCl box)
    + BLAKE2b-24 nonce derivation (stdlib hashlib)

Usage:
    kikemail.py generate [--domain kikemail.net]
    kikemail.py check <address|userpart> [--key HEX | --keyfile PATH]
                    [--host kikemail.net] [--raw] [--save DIR]

The private key never touches the network. Only the public key (as the
address userpart) is sent to the server when fetching sealed mail.
"""

import argparse
import base64
import email
import email.policy
import getpass
import hashlib
import json
import os
import secrets
import sys
import urllib.request
import urllib.error

# ---------------------------------------------------------------------------
# X25519 (RFC 7748) - Montgomery ladder
# ---------------------------------------------------------------------------

_P = 2**255 - 19
_A24 = 121665


def x25519(scalar: bytes, point: bytes) -> bytes:
    """X25519 scalar multiplication, RFC 7748 section 5."""
    k = bytearray(scalar)
    k[0] &= 248
    k[31] &= 127
    k[31] |= 64
    k_int = int.from_bytes(k, "little")
    u = int.from_bytes(point, "little") & ((1 << 255) - 1)

    x1 = u
    x2, z2 = 1, 0
    x3, z3 = u, 1
    swap = 0
    for t in range(254, -1, -1):
        kt = (k_int >> t) & 1
        swap ^= kt
        if swap:
            x2, x3 = x3, x2
            z2, z3 = z3, z2
        swap = kt
        a = (x2 + z2) % _P
        aa = (a * a) % _P
        b = (x2 - z2) % _P
        bb = (b * b) % _P
        e = (aa - bb) % _P
        c = (x3 + z3) % _P
        d = (x3 - z3) % _P
        da = (d * a) % _P
        cb = (c * b) % _P
        x3 = ((da + cb) ** 2) % _P
        z3 = (x1 * ((da - cb) ** 2 % _P)) % _P
        x2 = (aa * bb) % _P
        z2 = (e * (aa + _A24 * e % _P)) % _P
    if swap:
        x2, x3 = x3, x2
        z2, z3 = z3, z2
    return (x2 * pow(z2, _P - 2, _P) % _P).to_bytes(32, "little")


_BASEPOINT = bytes([9]) + bytes(31)


def public_from_secret(sk: bytes) -> bytes:
    return x25519(sk, _BASEPOINT)


# ---------------------------------------------------------------------------
# Salsa20 core / HSalsa20 / XSalsa20
# ---------------------------------------------------------------------------

_SIGMA = b"expand 32-byte k"


def _rotl32(v: int, n: int) -> int:
    return ((v << n) | (v >> (32 - n))) & 0xFFFFFFFF


def _quarterround(st: list, a: int, b: int, c: int, d: int) -> None:
    st[b] ^= _rotl32((st[a] + st[d]) & 0xFFFFFFFF, 7)
    st[c] ^= _rotl32((st[b] + st[a]) & 0xFFFFFFFF, 9)
    st[d] ^= _rotl32((st[c] + st[b]) & 0xFFFFFFFF, 13)
    st[a] ^= _rotl32((st[d] + st[c]) & 0xFFFFFFFF, 18)


def _salsa20_rounds(st: list) -> list:
    """20 rounds (10 double rounds). Returns transformed state (no feed-forward)."""
    w = st[:]
    for _ in range(10):
        # column round
        _quarterround(w, 0, 4, 8, 12)
        _quarterround(w, 5, 9, 13, 1)
        _quarterround(w, 10, 14, 2, 6)
        _quarterround(w, 15, 3, 7, 11)
        # row round
        _quarterround(w, 0, 1, 2, 3)
        _quarterround(w, 5, 6, 7, 4)
        _quarterround(w, 10, 11, 8, 9)
        _quarterround(w, 15, 12, 13, 14)
    return w


def _words(b: bytes) -> list:
    return [int.from_bytes(b[i:i + 4], "little") for i in range(0, len(b), 4)]


def hsalsa20(key: bytes, nonce16: bytes) -> bytes:
    """HSalsa20: key + 16-byte nonce -> 32-byte subkey (no feed-forward)."""
    c = _words(_SIGMA)
    k = _words(key)
    n = _words(nonce16)
    st = [c[0], k[0], k[1], k[2], k[3], c[1], n[0], n[1], n[2], n[3],
          c[2], k[4], k[5], k[6], k[7], c[3]]
    w = _salsa20_rounds(st)
    out = [w[0], w[5], w[10], w[15], w[6], w[7], w[8], w[9]]
    return b"".join(x.to_bytes(4, "little") for x in out)


def _salsa20_block(key: bytes, nonce8: bytes, counter: int) -> bytes:
    c = _words(_SIGMA)
    k = _words(key)
    n = _words(nonce8)
    st = [c[0], k[0], k[1], k[2], k[3], c[1], n[0], n[1],
          counter & 0xFFFFFFFF, (counter >> 32) & 0xFFFFFFFF,
          c[2], k[4], k[5], k[6], k[7], c[3]]
    w = _salsa20_rounds(st)
    # feed-forward
    return b"".join(((w[i] + st[i]) & 0xFFFFFFFF).to_bytes(4, "little")
                    for i in range(16))


def xsalsa20_stream(key: bytes, nonce24: bytes, length: int, start_block: int) -> bytes:
    subkey = hsalsa20(key, nonce24[:16])
    nonce8 = nonce24[16:]
    out = bytearray()
    block = start_block
    while len(out) < length:
        out += _salsa20_block(subkey, nonce8, block)
        block += 1
    return bytes(out[:length])


# ---------------------------------------------------------------------------
# Poly1305 (RFC 8439)
# ---------------------------------------------------------------------------


def poly1305(msg: bytes, key: bytes) -> bytes:
    r = int.from_bytes(key[:16], "little") & 0x0FFFFFFC0FFFFFFC0FFFFFFC0FFFFFFF
    s = int.from_bytes(key[16:], "little")
    p = (1 << 130) - 5
    acc = 0
    for i in range(0, len(msg), 16):
        block = msg[i:i + 16]
        # trailing 0x01 byte handles both full blocks (2^128) and partials
        n = int.from_bytes(block + b"\x01", "little")
        acc = ((acc + n) * r) % p
    acc = (acc + s) % (1 << 128)
    return acc.to_bytes(16, "little")


# ---------------------------------------------------------------------------
# NaCl box + sealed box (crypto_box_seal) open
# ---------------------------------------------------------------------------


def _box_before(pk: bytes, sk: bytes) -> bytes:
    """NaCl crypto_box_beforenm: X25519 then HSalsa20 with a zero nonce."""
    return hsalsa20(x25519(sk, pk), bytes(16))


def _box_open_afternm(ct: bytes, nonce24: bytes, shared: bytes) -> bytes:
    if len(ct) < 16:
        raise ValueError("ciphertext too short")
    mac, enc = ct[:16], ct[16:]
    # XSalsa20 keystream is continuous: bytes 0..31 are the Poly1305
    # one-time key, bytes 32+ XOR the message. So generate 32 + len(enc)
    # bytes from block 0 and split.
    stream = xsalsa20_stream(shared, nonce24, 32 + len(enc), 0)
    if not secrets.compare_digest(poly1305(enc, stream[:32]), mac):
        raise ValueError("decryption failed (bad key or corrupted data)")
    return bytes(a ^ b for a, b in zip(enc, stream[32:]))


def seal_open(sealed: bytes, pk: bytes, sk: bytes) -> bytes:
    """Open a libsodium/go crypto_box_seal (SealAnonymous) message."""
    if len(sealed) < 48:
        raise ValueError("sealed box too short")
    epk, ct = sealed[:32], sealed[32:]
    nonce = hashlib.blake2b(epk + pk, digest_size=24).digest()
    return _box_open_afternm(ct, nonce, _box_before(epk, sk))


# ---------------------------------------------------------------------------
# Base36 address encoding (matches server: left-pad to 32 bytes on decode)
# ---------------------------------------------------------------------------

_B36 = "0123456789abcdefghijklmnopqrstuvwxyz"


def b36_encode(raw: bytes) -> str:
    n = int.from_bytes(raw, "big")
    if n == 0:
        return "0"
    out = ""
    while n:
        n, r = divmod(n, 36)
        out = _B36[r] + out
    return out


def b36_decode(s: str) -> bytes:
    n = int(s, 36)
    raw = n.to_bytes(max(1, (n.bit_length() + 7) // 8), "big")
    if len(raw) > 32:
        raise ValueError("userpart decodes to more than 32 bytes - not a valid key")
    return raw.rjust(32, b"\x00")


# ---------------------------------------------------------------------------
# Mail rendering
# ---------------------------------------------------------------------------


def render_message(raw: bytes, raw_out: bool = False) -> str:
    """Render a decrypted payload. The server seals a JSON envelope
    {"body": <raw message>, "from": ..., "subject": ...}; the body itself
    is RFC822-ish (LF endings). Unwrap the JSON, then parse the body."""
    if raw_out:
        return raw.decode("utf-8", "replace")

    meta = {}
    try:
        env = json.loads(raw.decode("utf-8"))
        if isinstance(env, dict) and "body" in env:
            meta = env
            raw = env["body"].encode("utf-8")
    except (ValueError, UnicodeDecodeError):
        pass  # not JSON - treat as raw message

    # Normalize LF to CRLF so the email parser is happy.
    raw = raw.replace(b"\r\n", b"\n").replace(b"\n", b"\r\n")
    msg = email.message_from_bytes(raw, policy=email.policy.default)

    lines = [
        "From:    %s" % (meta.get("from") or msg.get("From", "")),
        "To:      %s" % msg.get("To", ""),
        "Subject: %s" % (meta.get("subject") or msg.get("Subject", "")),
        "Date:    %s" % msg.get("Date", ""),
        "-" * 60,
    ]
    bodies = []
    if msg.is_multipart():
        parts = [p for p in msg.walk() if p.get_content_type() == "text/plain"]
        if not parts:
            parts = [p for p in msg.walk() if p.get_content_maintype() == "text"]
        bodies = parts
    else:
        bodies = [msg]
    for part in bodies:
        payload = part.get_payload(decode=True)  # handles QP + base64
        if payload is None:
            continue
        charset = part.get_content_charset() or "utf-8"
        try:
            text = payload.decode(charset, "replace")
        except (LookupError, ValueError):
            text = payload.decode("utf-8", "replace")
        if part.get_content_type() != "text/plain":
            text = "[%s]\n%s" % (part.get_content_type(), text)
        lines.append(text.rstrip())
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def cmd_generate(args) -> int:
    sk = secrets.token_bytes(32)
    pk = public_from_secret(sk)
    user = b36_encode(pk)
    print("SAVE THIS DOWN - the private key is unrecoverable.\n")
    print("private key: %s" % sk.hex())
    print("address:     %s@%s" % (user, args.domain))
    return 0


def _fetch(host: str, user: str) -> dict:
    url = "https://%s/api/fetch/%s" % (host, user)
    req = urllib.request.Request(url, headers={"User-Agent": "kikemail-cli/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        raise SystemExit("fetch failed: HTTP %d from %s" % (e.code, url))
    except urllib.error.URLError as e:
        raise SystemExit("fetch failed: %s" % e.reason)


def cmd_check(args) -> int:
    addr = args.address.strip()
    user = addr.split("@")[0] if "@" in addr else addr
    host = args.host or (addr.split("@")[1] if "@" in addr else "kikemail.net")

    if args.keyfile:
        with open(args.keyfile) as f:
            key_hex = f.read().strip()
    elif args.key:
        key_hex = args.key.strip()
    else:
        key_hex = getpass.getpass("private key (hex): ").strip()
    if len(key_hex) != 64:
        raise SystemExit("invalid key length (expected 64 hex chars)")

    sk = bytes.fromhex(key_hex)
    try:
        pk = b36_decode(user.lower())
    except ValueError as e:
        raise SystemExit(str(e))

    # sanity: derived pubkey must match the address, else typo'd key
    derived = public_from_secret(sk)
    if derived != pk:
        print("warning: private key does not match this address", file=sys.stderr)

    resp = _fetch(host, user)
    rows = resp.get("mail") or []
    print("Address: %s@%s" % (user, host))
    print("Mail:    %d message(s)\n" % len(rows))

    if args.save:
        os.makedirs(args.save, exist_ok=True)

    ok, fail = 0, 0
    for i, row in enumerate(rows):
        try:
            pt = seal_open(bytes.fromhex(row["data"]), pk, sk)
        except ValueError as e:
            print("[message %d] %s" % (i + 1, e), file=sys.stderr)
            fail += 1
            continue
        ok += 1
        header = "=== message %d - received %s ===" % (i + 1, row.get("added", "?"))
        if args.save:
            path = os.path.join(args.save, "mail-%02d.eml" % (i + 1))
            with open(path, "wb") as f:
                f.write(pt)
            print("%s -> %s" % (header, path))
        else:
            print(header)
            print(render_message(pt, args.raw))
            print()
    print("%d decrypted, %d failed" % (ok, fail))
    return 0 if fail == 0 else 1


def main() -> int:
    ap = argparse.ArgumentParser(
        description="kikemail: addresses & encrypted inbox without the website")
    sub = ap.add_subparsers(dest="cmd", required=True)

    g = sub.add_parser("generate", help="generate a new keypair/address")
    g.add_argument("--domain", default="kikemail.net",
                   help="address domain (default: kikemail.net)")
    g.set_defaults(fn=cmd_generate)

    c = sub.add_parser("check", help="fetch and decrypt your inbox")
    c.add_argument("address", help="full address or just the userpart")
    c.add_argument("--key", help="private key hex (prompted if omitted; "
                                 "warning: visible in shell history)")
    c.add_argument("--keyfile", help="file containing private key hex")
    c.add_argument("--host", help="fetch host (default: from address domain)")
    c.add_argument("--raw", action="store_true", help="print raw RFC822 source")
    c.add_argument("--save", metavar="DIR", help="save decrypted .eml files to DIR")
    c.set_defaults(fn=cmd_check)

    args = ap.parse_args()
    return args.fn(args)


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