#!/usr/bin/env python3
"""Trace & Anchor desktop agent — drop a file in a folder, get a certificate back.

    python ta_agent.py setup        one-time: API key, folders, protected passphrase
    python ta_agent.py run          watch the folders (leave it running)
    python ta_agent.py anchor FILE [--protected] [--identity]
    python ta_agent.py verify FILE [--protected]
    python ta_agent.py status

Folders (created by `setup`, default on the Desktop):

    Trace & Anchor/
      1 - Standard (SHA-256)/   files with no names, addresses or personal data
      2 - Protected (HMAC)/     files that identify people (contracts, attestations)
      Certificates/             PDF certificate + proof JSON for every anchored file

Files never leave this computer. The agent computes the fingerprint locally and
sends only that 64-character digest to the Trace & Anchor API.

Protected mode: digest = HMAC-SHA-256(K, file), K = PBKDF2-HMAC-SHA256(passphrase,
salt, 600 000, 32 bytes) — byte-for-byte the scheme of the website, so a protected
anchor made here verifies on https://www.traceandanchor.com/verify-hash with the
passphrase. The passphrase is chosen once at setup and never stored or sent; only
the derived key K is kept, encrypted with Windows DPAPI (or a 0600 file elsewhere).

Standard library only (Python 3.9+), so it runs anywhere without installing anything.
"""

from __future__ import annotations

import argparse
import base64
import getpass
import hashlib
import hmac
import json
import logging
import os
import secrets
import sys
import time
import unicodedata
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from pathlib import Path

VERSION = '1.0.0'
DEFAULT_API = 'https://api-traceandanchor.com'
KDF_ITERATIONS = 600_000           # OWASP 2023 floor; must match src/lib/anchorDigest.ts
KDF_ALGORITHM = 'PBKDF2-HMAC-SHA256'
MIN_PASSPHRASE = 10
CHUNK = 1024 * 1024
SCAN_SECONDS = 5
STABLE_SECONDS = 3                 # a file must stop changing before we fingerprint it
BACKOFF_SECONDS = 300              # after "insufficient funds" or an outage

STANDARD_DIR = '1 - Standard (SHA-256)'
PROTECTED_DIR = '2 - Protected (HMAC)'
CERT_DIR = 'Certificates'
IGNORED_PREFIXES = ('~$', '.')
IGNORED_SUFFIXES = ('.tmp', '.part', '.crdownload', '.download', '.partial', '.lock')
IGNORED_NAMES = {'desktop.ini', 'thumbs.db', '.ds_store'}

CONFIG_DIR = Path(os.environ.get('TA_AGENT_HOME') or (Path.home() / '.traceandanchor'))
CONFIG_FILE = CONFIG_DIR / 'agent.json'
KEY_FILE = CONFIG_DIR / 'protected.key'

log = logging.getLogger('ta-agent')


# ── digests (must stay identical to the website) ─────────────────────────────

def derive_key(passphrase: str, salt_hex: str, iterations: int = KDF_ITERATIONS) -> bytes:
    """K = PBKDF2-HMAC-SHA256(NFKC(passphrase), salt, iterations, 32 bytes)."""
    normalized = unicodedata.normalize('NFKC', passphrase).encode('utf-8')
    return hashlib.pbkdf2_hmac('sha256', normalized, bytes.fromhex(salt_hex), iterations, 32)


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open('rb') as f:
        for chunk in iter(lambda: f.read(CHUNK), b''):
            h.update(chunk)
    return h.hexdigest()


def hmac_file(path: Path, key: bytes) -> str:
    h = hmac.new(key, digestmod=hashlib.sha256)
    with path.open('rb') as f:
        for chunk in iter(lambda: f.read(CHUNK), b''):
            h.update(chunk)
    return h.hexdigest()


# ── protected key storage ────────────────────────────────────────────────────

def _dpapi(data: bytes, protect: bool) -> bytes:
    import ctypes
    from ctypes import wintypes

    class Blob(ctypes.Structure):
        _fields_ = [('cbData', wintypes.DWORD), ('pbData', ctypes.POINTER(ctypes.c_char))]

    buf = ctypes.create_string_buffer(data, len(data))
    blob_in = Blob(len(data), ctypes.cast(buf, ctypes.POINTER(ctypes.c_char)))
    blob_out = Blob()
    fn = ctypes.windll.crypt32.CryptProtectData if protect else ctypes.windll.crypt32.CryptUnprotectData
    if protect:
        ok = fn(ctypes.byref(blob_in), 'Trace & Anchor agent', None, None, None, 0, ctypes.byref(blob_out))
    else:
        ok = fn(ctypes.byref(blob_in), None, None, None, None, 0, ctypes.byref(blob_out))
    if not ok:
        raise OSError('Windows DPAPI call failed.')
    try:
        return ctypes.string_at(blob_out.pbData, blob_out.cbData)
    finally:
        ctypes.windll.kernel32.LocalFree(blob_out.pbData)


def store_key(key: bytes) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    if os.name == 'nt':
        KEY_FILE.write_bytes(b'dpapi:' + base64.b64encode(_dpapi(key, True)))
    else:
        KEY_FILE.write_bytes(b'raw:' + base64.b64encode(key))
        os.chmod(KEY_FILE, 0o600)


def load_key() -> bytes | None:
    if not KEY_FILE.exists():
        return None
    blob = KEY_FILE.read_bytes()
    kind, _, payload = blob.partition(b':')
    raw = base64.b64decode(payload)
    return _dpapi(raw, False) if kind == b'dpapi' else raw


# ── config & state ───────────────────────────────────────────────────────────

def load_config() -> dict:
    if not CONFIG_FILE.exists():
        sys.exit('Not configured yet. Run:  python ta_agent.py setup')
    return json.loads(CONFIG_FILE.read_text(encoding='utf-8'))


def save_config(cfg: dict) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    CONFIG_FILE.write_text(json.dumps(cfg, indent=2), encoding='utf-8')
    if os.name != 'nt':
        os.chmod(CONFIG_FILE, 0o600)


class State:
    """Remembers what was anchored, so restarts never re-anchor (or re-charge)."""

    def __init__(self, root: Path):
        self.path = root / '.ta-agent' / 'state.json'
        self.path.parent.mkdir(parents=True, exist_ok=True)
        try:
            self.data = json.loads(self.path.read_text(encoding='utf-8'))
        except (OSError, ValueError):
            self.data = {'files': {}, 'digests': {}}

    def save(self):
        tmp = self.path.with_suffix('.tmp')
        tmp.write_text(json.dumps(self.data, indent=1), encoding='utf-8')
        tmp.replace(self.path)

    @staticmethod
    def signature(path: Path) -> str:
        st = path.stat()
        return f"{st.st_size}:{int(st.st_mtime)}"

    def seen(self, path: Path) -> bool:
        return self.data['files'].get(str(path)) == self.signature(path)

    def remember(self, path: Path, digest: str, anchor_id: str):
        self.data['files'][str(path)] = self.signature(path)
        self.data['digests'][digest] = anchor_id
        self.save()

    def mark_refused(self, path: Path):
        """Don't retry a file the API refused until it changes (new size/mtime)."""
        self.data['files'][str(path)] = self.signature(path)
        self.save()


# ── API client ───────────────────────────────────────────────────────────────

class ApiError(Exception):
    def __init__(self, status: int, code: str, detail: str):
        super().__init__(f"{status} {code}: {detail}")
        self.status, self.code, self.detail = status, code, detail


class Client:
    def __init__(self, base: str, key: str):
        self.base = base.rstrip('/')
        self.key = key

    def _request(self, method: str, path: str, body=None, headers=None, raw=False):
        data = json.dumps(body).encode('utf-8') if body is not None else None
        req = urllib.request.Request(self.base + path, data=data, method=method)
        req.add_header('Authorization', f'Bearer {self.key}')
        req.add_header('User-Agent', f'ta-agent/{VERSION}')
        if data is not None:
            req.add_header('Content-Type', 'application/json')
        for k, v in (headers or {}).items():
            req.add_header(k, v)
        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                payload = resp.read()
                return (payload, resp.headers) if raw else json.loads(payload or b'{}')
        except urllib.error.HTTPError as e:
            try:
                err = json.loads(e.read() or b'{}')
            except ValueError:
                err = {}
            raise ApiError(e.code, err.get('code', 'error'), err.get('detail', e.reason)) from None
        except urllib.error.URLError as e:
            raise ApiError(0, 'network_error', str(e.reason)) from None

    def account(self):
        return self._request('GET', '/api/v1/account')

    def find(self, digest: str):
        return self._request('GET', f'/api/v1/anchors?file_hash={digest}&limit=1')['data']

    def anchor(self, body: dict, idempotency_key: str):
        return self._request('POST', '/api/v1/anchors', body, {'Idempotency-Key': idempotency_key})

    def download(self, anchor_id: str, what: str) -> tuple[bytes, str]:
        payload, headers = self._request('GET', f'/api/v1/anchors/{anchor_id}/{what}', raw=True)
        return payload, _filename_from(headers.get('Content-Disposition', ''))

    def verify(self, digest: str):
        return self._request('GET', f'/api/v1/verify/{digest}')

    def new_identity_session(self):
        return self._request('POST', '/api/v1/identity/sessions', {})

    def identity_session(self, sid: str):
        return self._request('GET', f'/api/v1/identity/sessions/{sid}')


def _filename_from(disposition: str) -> str:
    for part in disposition.split(';'):
        part = part.strip()
        if part.lower().startswith("filename*=utf-8''"):
            return urllib.parse.unquote(part[17:])
    for part in disposition.split(';'):
        part = part.strip()
        if part.lower().startswith('filename='):
            return part[9:].strip('"')
    return ''


# ── anchoring a file ─────────────────────────────────────────────────────────

def anchor_file(client: Client, cfg: dict, path: Path, *, protected: bool,
                identity_session_id: str | None = None, out_dir: Path | None = None) -> dict:
    if protected:
        key = load_key()
        if key is None:
            raise SystemExit('Protected mode is not set up. Run setup again and choose a passphrase.')
        digest = hmac_file(path, key)
        body = {'hash_mode': 'hmac-sha256', 'kdf_salt': cfg['kdf_salt'],
                'kdf_iterations': cfg['kdf_iterations'], 'kdf_algorithm': KDF_ALGORITHM}
    else:
        digest = sha256_file(path)
        body = {'hash_mode': 'sha256'}

    existing = client.find(digest)
    if existing and not identity_session_id:
        result = existing[0]
        log.info('already anchored: %s -> %s', path.name, result['certificate_number'])
    else:
        body.update({'file_hash': digest, 'original_filename': path.name,
                     'original_size_bytes': path.stat().st_size})
        if identity_session_id:
            body['identity_session_id'] = identity_session_id
        # Same file + same mode -> same key: a retry after a timeout is replayed
        # by the server instead of charged twice.
        idem = hashlib.sha256(f"{body['hash_mode']}:{digest}:{identity_session_id or ''}".encode()).hexdigest()
        result = client.anchor(body, idem)
        log.info('anchored: %s -> %s (%s)', path.name, result['certificate_number'], result['level_label'])

    out_dir = out_dir or path.parent
    out_dir.mkdir(parents=True, exist_ok=True)
    for what, fallback in (('certificate.pdf', '.pdf'), ('proof.json', '-proof.json')):
        try:
            data, name = client.download(result['id'], what)
        except ApiError as e:
            log.warning('could not download %s for %s: %s', what, path.name, e)
            continue
        target = out_dir / (name or f"{path.stem}_{result['certificate_id']}{fallback}")
        target.write_bytes(data)
    return {'digest': digest, **result}


# ── commands ─────────────────────────────────────────────────────────────────

def cmd_setup(args):
    print(f'Trace & Anchor desktop agent {VERSION} — setup\n')
    cfg = json.loads(CONFIG_FILE.read_text(encoding='utf-8')) if CONFIG_FILE.exists() else {}
    api = input(f"API address [{cfg.get('api', DEFAULT_API)}]: ").strip() or cfg.get('api', DEFAULT_API)
    key = getpass.getpass('API key (ta_live_…, from Dashboard → Pro API): ').strip() or cfg.get('api_key', '')
    client = Client(api, key)
    try:
        acct = client.account()
    except ApiError as e:
        sys.exit(f'The API key was refused: {e.detail}')
    print(f"  ✓ connected as {acct['email']} — balance €{acct['wallet']['balance_cents'] / 100:.2f}")

    default_root = cfg.get('root') or str(Path.home() / 'Desktop' / 'Trace & Anchor')
    root = Path(input(f'Folder [{default_root}]: ').strip() or default_root).expanduser()
    for sub in (STANDARD_DIR, PROTECTED_DIR, CERT_DIR):
        (root / sub).mkdir(parents=True, exist_ok=True)
    cfg.update({'api': api, 'api_key': key, 'root': str(root)})

    if KEY_FILE.exists() and input('A protected passphrase is already set. Keep it? [Y/n]: ').strip().lower() != 'n':
        pass
    else:
        print('\nProtected mode (HMAC). Choose a passphrase ONCE. It is never sent to us')
        print('and cannot be recovered: without it, protected anchors can never be verified.')
        print('Leave empty to disable the Protected folder.\n')
        while True:
            p1 = getpass.getpass('Passphrase: ')
            if not p1:
                cfg.pop('kdf_salt', None)
                KEY_FILE.unlink(missing_ok=True)
                break
            if len(p1) < MIN_PASSPHRASE:
                print(f'  at least {MIN_PASSPHRASE} characters, please (four unrelated words is ideal)')
                continue
            if getpass.getpass('Again: ') != p1:
                print('  the two entries differ, try again')
                continue
            salt = secrets.token_hex(16)
            print('  deriving key (a few seconds)…')
            store_key(derive_key(p1, salt))
            cfg.update({'kdf_salt': salt, 'kdf_iterations': KDF_ITERATIONS})
            sheet = root / CERT_DIR / 'RECOVERY-SHEET-protected-mode.txt'
            if input(f'Write a recovery sheet with the passphrase to {sheet}? [y/N]: ').strip().lower() == 'y':
                sheet.write_text(_recovery_sheet(p1, salt), encoding='utf-8')
                print('  ✓ written — print it, store it safely, then delete the file')
            break

    save_config(cfg)
    print(f'\nDone. Folders are in {root}. Start watching with:  python ta_agent.py run')


def _recovery_sheet(passphrase: str, salt: str) -> str:
    return f"""TRACE & ANCHOR — DESKTOP AGENT — PROTECTED MODE RECOVERY SHEET
================================================================
WITHOUT THIS PASSPHRASE, FILES ANCHORED FROM THE "PROTECTED" FOLDER
CAN NEVER BE VERIFIED AGAIN — BY YOU, BY A COURT, OR BY US.

Passphrase          {passphrase}
Key derivation      {KDF_ALGORITHM}, {KDF_ITERATIONS} iterations
Salt (public)       {salt}

To verify a file: https://www.traceandanchor.com/verify-hash
select the original file (it is not uploaded) and enter the passphrase.
================================================================
"""


def _usable(path: Path, root: Path) -> bool:
    if not path.is_file():
        return False
    rel = path.relative_to(root).parts
    if any(p.startswith('.') for p in rel) or CERT_DIR in rel:
        return False
    name = path.name.lower()
    return not (name in IGNORED_NAMES or path.name.startswith(IGNORED_PREFIXES) or name.endswith(IGNORED_SUFFIXES))


def cmd_run(args):
    cfg = load_config()
    root = Path(cfg['root'])
    client = Client(cfg['api'], cfg['api_key'])
    state = State(root)
    certs = root / CERT_DIR
    pending: dict[str, tuple[str, float]] = {}
    paused_until = 0.0
    log.info('watching %s (Ctrl+C to stop)', root)

    while True:
        if time.time() < paused_until:
            time.sleep(SCAN_SECONDS)
            continue
        for folder, protected in ((root / STANDARD_DIR, False), (root / PROTECTED_DIR, True)):
            if protected and not cfg.get('kdf_salt'):
                continue
            for path in sorted(folder.rglob('*')):
                try:
                    if not _usable(path, root) or state.seen(path):
                        continue
                    sig = State.signature(path)
                    first = pending.get(str(path))
                    if first is None or first[0] != sig:
                        pending[str(path)] = (sig, time.time())
                        continue
                    if time.time() - first[1] < STABLE_SECONDS:
                        continue
                    result = anchor_file(client, cfg, path, protected=protected, out_dir=certs)
                    state.remember(path, result['digest'], result['id'])
                    pending.pop(str(path), None)
                except ApiError as e:
                    if e.code in ('insufficient_funds', 'network_error', 'service_unavailable', 'upstream_error', 'rate_limited'):
                        log.error('%s — pausing %d min: %s', path.name, BACKOFF_SECONDS // 60, e.detail)
                        paused_until = time.time() + BACKOFF_SECONDS
                        break
                    log.error('%s refused: %s', path.name, e.detail)
                    state.mark_refused(path)
                except OSError as e:
                    log.warning('%s not readable yet: %s', path.name, e)
        time.sleep(SCAN_SECONDS)


def _wait_for_identity(client: Client) -> str:
    s = client.new_identity_session()
    print('\nIdentity-level anchor: approve with FranceConnect in your browser:')
    print(f"  {s['approval_url']}\n")
    webbrowser.open(s['approval_url'])
    while True:
        time.sleep(3)
        cur = client.identity_session(s['id'])
        if cur['status'] == 'approved':
            print('  ✓ approved')
            return s['id']
        if cur['status'] in ('denied', 'expired'):
            sys.exit(f"Identity session {cur['status']}.")


def cmd_anchor(args):
    cfg = load_config()
    client = Client(cfg['api'], cfg['api_key'])
    sid = _wait_for_identity(client) if args.identity else None
    result = anchor_file(client, cfg, Path(args.file), protected=args.protected,
                         identity_session_id=sid, out_dir=Path(cfg['root']) / CERT_DIR)
    print(json.dumps({k: result[k] for k in ('id', 'certificate_number', 'level_label', 'file_hash')}, indent=2))


def cmd_verify(args):
    cfg = load_config()
    client = Client(cfg['api'], cfg['api_key'])
    path = Path(args.file)
    digest = hmac_file(path, load_key()) if args.protected else sha256_file(path)
    res = client.verify(digest)
    if res.get('found'):
        print(f'✓ anchored — {digest}')
        for r in res.get('database_records', []):
            print(f"  {r['anchored_at']}  {r.get('solana_tx_url') or r.get('solana_signature')}")
    else:
        print(f'✗ no anchor found for {digest}')
        sys.exit(1)


def cmd_status(args):
    cfg = load_config()
    acct = Client(cfg['api'], cfg['api_key']).account()
    print(f"account   {acct['email']}\nbalance   €{acct['wallet']['balance_cents'] / 100:.2f}")
    print(f"folder    {cfg['root']}\nprotected {'on' if cfg.get('kdf_salt') else 'off'}")


def main(argv=None):
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
    p = argparse.ArgumentParser(prog='ta_agent', description='Trace & Anchor desktop agent')
    sub = p.add_subparsers(dest='cmd', required=True)
    sub.add_parser('setup').set_defaults(fn=cmd_setup)
    sub.add_parser('run').set_defaults(fn=cmd_run)
    sub.add_parser('status').set_defaults(fn=cmd_status)
    a = sub.add_parser('anchor')
    a.add_argument('file')
    a.add_argument('--protected', action='store_true')
    a.add_argument('--identity', action='store_true', help='identity-level anchor (FranceConnect approval)')
    a.set_defaults(fn=cmd_anchor)
    v = sub.add_parser('verify')
    v.add_argument('file')
    v.add_argument('--protected', action='store_true')
    v.set_defaults(fn=cmd_verify)
    args = p.parse_args(argv)
    try:
        args.fn(args)
    except KeyboardInterrupt:
        print('\nstopped')


if __name__ == '__main__':
    main()
