#!/usr/bin/env python3
"""AirNDB remote packet capture helper over SSH."""
from __future__ import annotations
import argparse, base64, json, os, shlex, shutil, subprocess
from datetime import datetime, timezone
from pathlib import Path
MAX_COUNT = 500000

def stamp(): return datetime.now().strftime('%Y%m%d-%H%M%S')
def iso(): return datetime.now(timezone.utc).isoformat()
def q(s): return shlex.quote(str(s))
def split(s): return shlex.split(s, posix=True) if s.strip() else []

def read_env(path):
    vals = {}
    if path.exists():
        for raw in path.read_text(encoding='utf-8-sig').splitlines():
            line = raw.strip()
            if line and not line.startswith('#') and '=' in line:
                k, v = line.split('=', 1); vals[k.strip()] = v.strip().strip('"').strip("'")
    for k in ['AIRNDB_REMOTE_SSH_TARGET','AIRNDB_REMOTE_SSH_PORT','AIRNDB_REMOTE_SSH_OPTIONS','AIRNDB_REMOTE_WORKDIR','AIRNDB_REMOTE_TCPDUMP','AIRNDB_REMOTE_CAPTURE_PREFIX']:
        if os.environ.get(k): vals[k] = os.environ[k]
    vals.setdefault('AIRNDB_REMOTE_SSH_PORT','22'); vals.setdefault('AIRNDB_REMOTE_TCPDUMP','auto'); vals.setdefault('AIRNDB_REMOTE_CAPTURE_PREFIX','sudo -n')
    return vals

def write_env(path, updates):
    path.parent.mkdir(parents=True, exist_ok=True); old = path.read_text(encoding='utf-8') if path.exists() else ''; keys = set(updates); lines = []
    for raw in old.splitlines():
        key = raw.split('=',1)[0].strip() if '=' in raw and not raw.strip().startswith('#') else None
        if key not in keys: lines.append(raw)
    if lines and lines[-1].strip(): lines.append('')
    for k, v in updates.items():
        if v: lines.append(f'{k}={v}')
    path.write_text('\n'.join(lines).rstrip()+'\n', encoding='utf-8', newline='\n')

def ensure_files(project):
    d = project/'AirPlan'/'state'/'airndb'; d.mkdir(parents=True, exist_ok=True); ex = d/'remote-device.env.example'
    if not ex.exists():
        ex.write_text('# AirNDB remote packet-capture device configuration\nAIRNDB_REMOTE_SSH_TARGET=user@host\nAIRNDB_REMOTE_SSH_PORT=22\nAIRNDB_REMOTE_SSH_OPTIONS=\nAIRNDB_REMOTE_WORKDIR=\nAIRNDB_REMOTE_TCPDUMP=auto\nAIRNDB_REMOTE_CAPTURE_PREFIX=sudo -n\n', encoding='utf-8', newline='\n')
    gi = d/'.gitignore'; old = gi.read_text(encoding='utf-8') if gi.exists() else ''
    for item in ['remote-device.env','tool.env']:
        if item not in old.splitlines(): old = (old.rstrip()+f'\n{item}\n').lstrip()
    gi.write_text(old, encoding='utf-8', newline='\n')

def ssh_args(env, cmd):
    target = env.get('AIRNDB_REMOTE_SSH_TARGET','')
    if not target: raise SystemExit('AIRNDB_REMOTE_SSH_TARGET is required')
    args = [shutil.which('ssh') or 'ssh']; port = env.get('AIRNDB_REMOTE_SSH_PORT','22')
    if port: args += ['-p', port]
    args += split(env.get('AIRNDB_REMOTE_SSH_OPTIONS','')); return args + [target, cmd]

def run(env, cmd, timeout=60):
    return subprocess.run(ssh_args(env, cmd), capture_output=True, text=True, encoding='utf-8', errors='replace', timeout=timeout)

def remote_home(env):
    r = run(env, 'printf %s "$HOME"', 15); return r.stdout.strip() if r.returncode == 0 else ''

def workdir(env):
    if env.get('AIRNDB_REMOTE_WORKDIR'): return env['AIRNDB_REMOTE_WORKDIR']
    home = remote_home(env); return (home.rstrip('/')+'/.airndb') if home else '.airndb'

def detect_tool(env):
    tool = env.get('AIRNDB_REMOTE_TCPDUMP','auto').strip()
    if tool and tool != 'auto': return tool
    r = run(env, 'for t in tcpdump dumpcap windump WinDump.exe; do command -v "$t" >/dev/null 2>&1 && { command -v "$t"; exit 0; }; done', 15)
    return r.stdout.strip().splitlines()[-1] if r.returncode == 0 and r.stdout.strip() else ''

def install_tool(env):
    script = """set -e
if command -v tcpdump >/dev/null 2>&1; then exit 0; fi
if [ "$(id -u 2>/dev/null || echo 1)" = "0" ]; then SUDO=""; else SUDO="sudo -n"; fi
if command -v apt-get >/dev/null 2>&1; then $SUDO apt-get update && $SUDO apt-get install -y tcpdump;
elif command -v dnf >/dev/null 2>&1; then $SUDO dnf install -y tcpdump;
elif command -v yum >/dev/null 2>&1; then $SUDO yum install -y tcpdump;
elif command -v apk >/dev/null 2>&1; then $SUDO apk add tcpdump;
elif command -v pacman >/dev/null 2>&1; then $SUDO pacman -Sy --noconfirm tcpdump;
else exit 42; fi"""
    return run(env, script, 300)

def setup(project, env, auto=True):
    if not env.get('AIRNDB_REMOTE_SSH_TARGET'):
        return {'status':'blocked','reason':'missing_remote_target','example':str(project/'AirPlan'/'state'/'airndb'/'remote-device.env.example')}
    if not shutil.which('ssh'): return {'status':'blocked','reason':'ssh_not_found'}
    probe = run(env, 'printf ok', 20)
    if probe.returncode: return {'status':'blocked','reason':'ssh_probe_failed','stderr':probe.stderr.strip()}
    wd = workdir(env); mk = run(env, f'mkdir -p {q(wd)}', 20)
    if mk.returncode: return {'status':'blocked','reason':'remote_workdir_failed','stderr':mk.stderr.strip()}
    tool = detect_tool(env); install = 'skipped'
    if not tool and auto:
        inst = install_tool(env); install = 'ok' if inst.returncode == 0 else f'failed:{inst.returncode}'; tool = detect_tool(env)
    upd = {'AIRNDB_REMOTE_WORKDIR':wd}
    if tool: upd['AIRNDB_REMOTE_TCPDUMP'] = tool
    write_env(project/'AirPlan'/'state'/'airndb'/'remote-device.env', upd)
    return {'status':'ok' if tool else 'blocked','target':env.get('AIRNDB_REMOTE_SSH_TARGET',''),'workdir':wd,'tool':tool,'capturePrefix':env.get('AIRNDB_REMOTE_CAPTURE_PREFIX',''),'autoConfigure':install,'hint':'' if tool else 'Install tcpdump/dumpcap or set AIRNDB_REMOTE_TCPDUMP.'}

def bounded(n):
    if n < 1 or n > MAX_COUNT: raise SystemExit(f'--count must be 1..{MAX_COUNT}')
    return n

def filt(expr):
    if not expr.strip(): return []
    try: return split(expr)
    except ValueError: return expr.split()

def cap_cmd(env, tool, iface, out, count, filter_expr):
    if not iface: raise SystemExit('--iface is required')
    parts = split(env.get('AIRNDB_REMOTE_CAPTURE_PREFIX','')) + [tool,'-i',iface,'-nn','-s','0','-w',out,'-c',str(bounded(count))] + filt(filter_expr)
    return shlex.join(parts)

def outdir(project, configured):
    p = Path(configured).expanduser().resolve() if configured else project/'AirPlan'/'docs'/'network'/'airndb-captures'; p.mkdir(parents=True, exist_ok=True); return p

def emit(prefix, data):
    print(f'{prefix}_status={data.get("status","unknown")}')
    for k, v in data.items():
        if k != 'status': print(f'{k}={json.dumps(v, ensure_ascii=False) if isinstance(v,(dict,list)) else v}')

def log(project, data):
    lp = project/'AirPlan'/'docs'/'network'/'airndb-log.md'; lp.parent.mkdir(parents=True, exist_ok=True)
    lp.open('a', encoding='utf-8', newline='\n').write(f"\n## {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}: AirNDB remote {data.get('action','capture')}\n\n- Remote target: `{data.get('target','')}`\n- Tool: `{data.get('tool','')}`\n- Interface: `{data.get('interface','')}`\n- Filter: `{data.get('filter','')}`\n- Command: `{data.get('command','')}`\n- Artifacts: `{data.get('artifact','')}`\n- Report: `{data.get('report','')}`\n- Result: `{data.get('status','')}`\n- AirDbg handoff: TODO\n- Residual risk: remote pcap files may contain sensitive packet contents.\n")

def main():
    ap = argparse.ArgumentParser(); ap.add_argument('--project', default='.'); ap.add_argument('--action', choices=['setup','status','interfaces','command','capture'], default='setup'); ap.add_argument('--iface', default=''); ap.add_argument('--filter', default=''); ap.add_argument('--count', type=int, default=200); ap.add_argument('--timeout', type=int, default=30); ap.add_argument('--output-dir', default=''); ap.add_argument('--output', default=''); ap.add_argument('--no-auto-configure', action='store_true'); a = ap.parse_args()
    project = Path(a.project).expanduser().resolve(); ensure_files(project); env = read_env(project/'AirPlan'/'state'/'airndb'/'remote-device.env')
    info = setup(project, env, auto=(a.action!='status' and not a.no_auto_configure))
    if a.action in ['setup','status'] or info.get('status') != 'ok': emit('airndb_remote', info); raise SystemExit(0 if info.get('status')=='ok' else 2)
    od = outdir(project, a.output_dir); tool = info['tool']
    if a.action == 'interfaces':
        icmd = shlex.join(split(env.get('AIRNDB_REMOTE_CAPTURE_PREFIX','')) + [tool,'-D']); r = run(env, icmd, max(5,a.timeout)); data = {'action':'interfaces','status':'ok' if r.returncode==0 else 'failed','target':info['target'],'tool':tool,'command':shlex.join(ssh_args(env, icmd)),'stdout':r.stdout.strip(),'stderr':r.stderr.strip()}; report = od/(stamp()+'-airndb-remote-interfaces.json'); report.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n', encoding='utf-8'); data['report']=str(report); log(project,data); emit('airndb_remote',data); print(r.stdout); print(r.stderr); raise SystemExit(0 if data['status']=='ok' else 1)
    remote = a.output or info['workdir'].rstrip('/') + '/' + stamp() + '-airndb-remote.pcap'; cmd = cap_cmd(env, tool, a.iface, remote, a.count, a.filter)
    if a.action == 'command': emit('airndb_remote', {'status':'ok','target':info['target'],'remotePcap':remote,'command':shlex.join(ssh_args(env, cmd))}); return
    r = run(env, cmd, max(1,a.timeout)); data = {'action':'capture','status':'ok' if r.returncode==0 else 'failed','target':info['target'],'tool':tool,'interface':a.iface,'filter':a.filter,'count':bounded(a.count),'timeoutSeconds':a.timeout,'remotePcap':remote,'command':shlex.join(ssh_args(env, cmd)),'stdout':r.stdout.strip(),'stderr':r.stderr.strip(),'capturedAt':iso()}
    if r.returncode == 0:
        b64 = run(env, f'base64 < {q(remote)}', a.timeout); local = od/(stamp()+'-airndb-remote.pcap'); local.write_bytes(base64.b64decode(''.join(b64.stdout.split()))); data['artifact'] = str(local)
    report = od/(stamp()+'-airndb-remote-capture.json'); report.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n', encoding='utf-8'); data['report'] = str(report); log(project, data); emit('airndb_remote', data); raise SystemExit(0 if data['status']=='ok' else 1)
if __name__ == '__main__': main()