#!/usr/bin/env python3
"""
HTTPS exec server for operator remote command execution.
Runs on bl (stable), accepts authenticated POST /exec, returns command output.

Usage:
    python3 exec-server.py --port 8443 --token-file /home/super/.exec-token

646 (or any operator) can then:
    curl -k -X POST https://100.123.153.75:8443/exec \
        -H "Content-Type: application/json" \
        -d '{"cmd": "dm.py send --agent 646 --to opm --target main \"hello\"", "token": "..."}'

Via the VM Caddy (no tailnet needed from the agent's container):
    curl -k -X POST https://34-139-37-135.sslip.io/exec/exec \
        -H "Content-Type: application/json" \
        -d '{"cmd": "...", "token": "<per-agent token>"}'

Signature auth (preferred for agents): no secret crosses the wire at all.
The agent signs {"cmd","ts","nonce"} with their registered SSH key
(`ssh-keygen -Y sign -n exec-server`) and posts
{"identity","payload","signature"}. Verified against the signers file
with `ssh-keygen -Y verify`; ts must be within 300s and the nonce unused.
This is the same identity primitive as signed board posts, and it never
trips secret-handling guardrails because a signature is not a secret.

Auth: the master token (TOKEN_FILE) plus per-agent tokens, one file per
agent under TOKEN_DIR (0600). Each agent's token is individually revocable
by deleting its file. The using identity is logged (never the token).
"""
import argparse
import hashlib
import hmac
import json
import secrets
import subprocess
import sys
from http.server import HTTPServer, BaseHTTPRequestHandler
import ssl

# Token file - generated on first run if not exists
TOKEN_FILE = '/home/super/.exec-server-token'
# Per-agent tokens: TOKEN_DIR/<agent> contains that agent's token
TOKEN_DIR = '/home/super/.exec-tokens'
# ssh-keygen signature auth: public keys of fleet identities, one per line
# ("<identity> ssh-ed25519 AAAA..."), synced from the VM's allowed_signers.
SIGNERS_FILE = '/home/super/.exec-signers'
# Seen nonces for replay protection ("<nonce> <ts>" per line).
NONCE_FILE = '/home/super/.exec-nonces'
SIG_NAMESPACE = 'exec-server'
SIG_MAX_SKEW = 300  # seconds; nonces remembered for 2x this

def get_token():
    """Load or generate the auth token."""
    try:
        with open(TOKEN_FILE, 'r') as f:
            return f.read().strip()
    except FileNotFoundError:
        token = secrets.token_hex(32)
        with open(TOKEN_FILE, 'w') as f:
            f.write(token)
        # Secure permissions
        import os
        os.chmod(TOKEN_FILE, 0o600)
        print(f'Generated new token in {TOKEN_FILE}', file=sys.stderr)
        return token

def check_token(token):
    """Check token against the master token and per-agent tokens.
    Returns the identity label ('master' or the agent filename), or None."""
    if token and hmac.compare_digest(token, get_token()):
        return 'master'
    import os
    try:
        names = os.listdir(TOKEN_DIR)
    except FileNotFoundError:
        return None
    for name in names:
        p = os.path.join(TOKEN_DIR, name)
        if not os.path.isfile(p):
            continue
        try:
            with open(p) as f:
                t = f.read().strip()
        except OSError:
            continue
        if t and hmac.compare_digest(token, t):
            return name
    return None

def check_nonce(nonce):
    """True if the nonce was never used; records it. Prunes expired entries."""
    import os
    import time
    now = time.time()
    fresh = []
    try:
        with open(NONCE_FILE) as f:
            for line in f:
                parts = line.split()
                if len(parts) != 2:
                    continue
                n, t = parts
                try:
                    if now - float(t) < 2 * SIG_MAX_SKEW:
                        fresh.append((n, t))
                except ValueError:
                    pass
    except FileNotFoundError:
        pass
    if any(n == nonce for n, _ in fresh):
        return False
    fresh.append((nonce, str(now)))
    try:
        with open(NONCE_FILE, 'w') as f:
            for n, t in fresh:
                f.write(f"{n} {t}\n")
        os.chmod(NONCE_FILE, 0o600)
    except OSError:
        return False
    return True


def check_signature(identity, payload, signature):
    """Verify an `ssh-keygen -Y` signature over the payload envelope.
    Returns the identity on success, None on failure. The envelope must be
    JSON {"cmd","ts","nonce"} with a fresh ts and an unused nonce."""
    import json as _json
    import os
    import re
    import subprocess
    import tempfile
    import time
    if not re.fullmatch(r'[a-z0-9-]+', identity or ''):
        return None
    try:
        data = _json.loads(payload)
    except Exception:
        return None
    if not isinstance(data, dict):
        return None
    cmd = data.get('cmd')
    ts = data.get('ts')
    nonce = data.get('nonce')
    if not isinstance(cmd, str) or not cmd:
        return None
    if not isinstance(nonce, str) or not re.fullmatch(r'[0-9a-fA-F]{16,128}', nonce):
        return None
    try:
        ts = float(ts)
    except (TypeError, ValueError):
        return None
    if abs(time.time() - ts) > SIG_MAX_SKEW:
        return None
    if not check_nonce(nonce):
        return None
    sig_path = None
    try:
        with tempfile.NamedTemporaryFile('w', delete=False, suffix='.sig') as f:
            f.write(signature if signature.endswith('\n') else signature + '\n')
            sig_path = f.name
        p = subprocess.run(
            ['ssh-keygen', '-Y', 'verify', '-f', SIGNERS_FILE, '-I', identity,
             '-n', SIG_NAMESPACE, '-s', sig_path],
            input=payload.encode(), capture_output=True, timeout=15)
        return identity if p.returncode == 0 else None
    except Exception:
        return None
    finally:
        if sig_path:
            try:
                os.unlink(sig_path)
            except OSError:
                pass


class ExecHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        # Quiet logging, just to stderr
        sys.stderr.write(f'{self.client_address[0]} - {format % args}\n')
    
    def do_POST(self):
        # Read body
        content_length = int(self.headers.get('Content-Length', 0))
        if content_length > 1024 * 1024:  # 1MB max
            self.send_response(413)
            self.end_headers()
            return

        body = self.rfile.read(content_length)
        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'{"error": "invalid json"}')
            return

        if self.path == '/exec/rotate':
            self.handle_rotate(data)
            return

        if self.path != '/exec':
            self.send_response(404)
            self.end_headers()
            return
        
        # Auth: bearer token OR ssh-keygen -Y signature (no secret in transit).
        # Bearer: {"token": "...", "cmd": "..."}.
        # Signature: {"identity": "...", "payload": "{\"cmd\":...,\"ts\":...,\"nonce\":...}",
        #             "signature": "<ssh-keygen -Y armor>"}. check_signature
        # validates the envelope; cmd comes from the signed payload.
        ident = None
        cmd = ''
        token = data.get('token', '')
        if token:
            ident = check_token(token)
            cmd = data.get('cmd', '')
        elif data.get('identity') and data.get('payload') and data.get('signature'):
            ident = check_signature(data['identity'], data['payload'], data['signature'])
            if ident:
                try:
                    cmd = json.loads(data['payload']).get('cmd', '')
                except Exception:
                    cmd = ''
        if not ident:
            self.send_response(401)
            self.end_headers()
            self.wfile.write(b'{"error": "unauthorized"}')
            return
        sys.stderr.write(f'exec as {ident} from {self.client_address[0]}\n')
        # Get command
        if not cmd or not isinstance(cmd, str):
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'{"error": "missing cmd"}')
            return
        
        # Execute (with timeout)
        timeout = min(data.get('timeout', 60), 300)  # max 5 min
        try:
            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd='/home/super'
            )
            response = {
                'stdout': result.stdout,
                'stderr': result.stderr,
                'rc': result.returncode,
            }
        except subprocess.TimeoutExpired:
            response = {'error': 'timeout', 'rc': -1}
        except Exception as e:
            response = {'error': str(e), 'rc': -1}
        
        # Send response
        resp_body = json.dumps(response).encode()
        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        self.send_header('Content-Length', str(len(resp_body)))
        self.end_headers()
        self.wfile.write(resp_body)
    
    def handle_rotate(self, data):
        """POST /exec/rotate - replace an agent's token with a fresh one.
        The old token dies immediately; the new one is returned ONLY in the
        response body (never logged). Lets an agent bootstrap from a
        trust-root-delivered token and end up with one nobody else knows.
        Agents may rotate only their own token; master may rotate any
        agent's (not its own - that stays manual on bl)."""
        import os
        import re
        token = data.get('token', '')
        ident = check_token(token)
        if not ident:
            self.send_response(401)
            self.end_headers()
            self.wfile.write(b'{"error": "unauthorized"}')
            return
        target = data.get('agent', ident)
        if not re.fullmatch(r'[a-z0-9-]+', target or ''):
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'{"error": "bad agent name"}')
            return
        if target == 'master' or (ident != 'master' and target != ident):
            self.send_response(403)
            self.end_headers()
            self.wfile.write(b'{"error": "forbidden"}')
            return
        path = os.path.join(TOKEN_DIR, target)
        if not os.path.isfile(path):
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b'{"error": "no such agent token"}')
            return
        new_token = secrets.token_hex(32)
        tmp = path + '.tmp'
        with open(tmp, 'w') as f:
            f.write(new_token + '\n')
        os.chmod(tmp, 0o600)
        os.replace(tmp, path)
        sys.stderr.write(f'token rotated for {target} by {ident}\n')
        resp_body = json.dumps({'token': new_token}).encode()
        self.send_response(200)
        self.send_header('Content-Type', 'application/json')
        self.send_header('Content-Length', str(len(resp_body)))
        self.end_headers()
        self.wfile.write(resp_body)

    def do_GET(self):
        if self.path == '/health':
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            self.wfile.write(b'{"status": "ok"}')
        else:
            self.send_response(404)
            self.end_headers()

def main():
    global TOKEN_FILE, TOKEN_DIR
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8443)
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--token-file', default='/home/super/.exec-server-token')
    parser.add_argument('--token-dir', default='/home/super/.exec-tokens')
    parser.add_argument('--signers-file', default='/home/super/.exec-signers')
    parser.add_argument('--nonce-file', default='/home/super/.exec-nonces')
    args = parser.parse_args()

    TOKEN_FILE = args.token_file
    TOKEN_DIR = args.token_dir
    global SIGNERS_FILE, NONCE_FILE
    SIGNERS_FILE = args.signers_file
    NONCE_FILE = args.nonce_file
    
    # Ensure token exists
    token = get_token()
    print(f'Token: {token[:8]}... (full in {TOKEN_FILE})', file=sys.stderr)
    
    server = HTTPServer((args.host, args.port), ExecHandler)
    
    # Wrap with TLS (self-signed is fine for our use, we use -k)
    # Generate self-signed cert if not exists
    import os
    cert_file = '/home/super/.exec-server-cert.pem'
    key_file = '/home/super/.exec-server-key.pem'
    
    if not os.path.exists(cert_file):
        print('Generating self-signed cert...', file=sys.stderr)
        subprocess.run([
            'openssl', 'req', '-x509', '-newkey', 'rsa:2048',
            '-keyout', key_file, '-out', cert_file,
            '-days', '3650', '-nodes',
            '-subj', '/CN=bl-exec-server'
        ], check=True, capture_output=True)
        os.chmod(key_file, 0o600)
    
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.load_cert_chain(cert_file, key_file)
    server.socket = context.wrap_socket(server.socket, server_side=True)
    
    print(f'Exec server listening on https://{args.host}:{args.port}/exec', file=sys.stderr)
    print('Health check: https://<host>:<port>/health', file=sys.stderr)
    
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print('\nShutting down', file=sys.stderr)

if __name__ == '__main__':
    main()