#!/usr/bin/env python3
"""
geo-scout — geo-targeted network scanner
masscan for port discovery, nmap for fingerprinting, probe engine for matching
"""

import ipaddress
import json
import os
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import yaml

BASE_DIR = Path(__file__).parent
sys.path.insert(0, str(BASE_DIR))

from lib.cidr import fetch_rir_data, load_index, get_cidrs, merge_cidrs
from lib.scanner import check_deps, masscan_sweep, nmap_fingerprint
from lib.prober import run_probe
from lib.matcher import match_target
from lib import db

INPUTS_DIR = BASE_DIR / "inputs"
DEFAULT_COUNTRIES = INPUTS_DIR / "countries.yaml"
DEFAULT_TARGETS = INPUTS_DIR / "targets.yaml"


# ── ANSI colors ───────────────────────────────────────────────────────────────
R = "\033[0;31m"
G = "\033[0;32m"
Y = "\033[0;33m"
C = "\033[0;36m"
B = "\033[1;34m"
W = "\033[0;37m"
BOLD = "\033[1m"
RST = "\033[0m"

def err(msg): print(f"{R}[!]{RST} {msg}", file=sys.stderr)
def ok(msg):  print(f"{G}[+]{RST} {msg}")
def info(msg): print(f"{C}[*]{RST} {msg}")
def warn(msg): print(f"{Y}[-]{RST} {msg}")
def banner(): print(f"""{B}
        ·    *    ·    *    ·    *    ·    *    ·
      *   ╔══════════════════════════════════╗   *
     ·    ║  ·  ·   ·   ·   ·   ·   ·  ·  ·║    ·
      *   ║ · ─────────────────────────── · ║   *
     ·    ║  ─────  ·    ╔═══╗    ·  ─────  ║    ·
      *   ║ · ──── · ────║{R} ◉ {B}║──── · ───· ║   *
     ·    ║  ─────  ·    ╚═══╝    ·  ─────  ║    ·
      *   ║ · ─────────────────────────── · ║   *
     ·    ║  ·  ·   ·   ·   ·   ·   ·  ·  ·║    ·
      *   ╚══════════════════════════════════╝   *
        ·    *    ·    *    ·    *    ·    *    ·
{RST}{C}
  ██████  ███████  ██████        ███████  ██████  ██████  ██    ██ ████████
 ██       ██      ██    ██       ██      ██      ██    ██ ██    ██    ██
 ██   ███ █████   ██    ██ ████  ███████ ██      ██    ██ ██    ██    ██
 ██    ██ ██      ██    ██            ██ ██      ██    ██ ██    ██    ██
  ██████  ███████  ██████        ███████  ██████  ██████   ██████     ██
{RST}{B}  ───────────────────────────────────────────────────────────────────────
{RST}{W}  geo-targeted threat surface scanner  {Y}v1.0{RST}  {R}[ authorized use only ]{RST}
""")


# ── Loaders ───────────────────────────────────────────────────────────────────
def load_countries(path: Path) -> dict:
    if not path.exists():
        err(f"Countries file not found: {path}")
        sys.exit(1)
    with open(path) as f:
        return yaml.safe_load(f)


def load_targets(path: Path) -> list[dict]:
    if not path.exists():
        err(f"Targets file not found: {path}")
        sys.exit(1)
    with open(path) as f:
        data = yaml.safe_load(f)
    return data.get("targets", [])


def collect_all_ports(targets: list[dict]) -> list[int]:
    ports = set()
    for t in targets:
        ports.update(t.get("ports", []))
    return sorted(ports)


def count_ips(cidrs: list[str]) -> int:
    total = 0
    for c in cidrs:
        try:
            total += ipaddress.ip_network(c, strict=False).num_addresses
        except ValueError:
            pass
    return total


def fmt_ip_count(n: int) -> str:
    if n >= 1_000_000:
        return f"{n/1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n/1_000:.0f}K"
    return str(n)


def fmt_eta(total_ips: int, n_ports: int, rate: int) -> str:
    if rate <= 0:
        return "?"
    secs = int(total_ips * n_ports / rate)
    m, s = divmod(secs, 60)
    h, m = divmod(m, 60)
    return f"{h}h {m:02d}m" if h else f"{m}m {s:02d}s"


# ── Probe worker ──────────────────────────────────────────────────────────────
def probe_host(ip: str, port: int, targets: list[dict]) -> list[dict]:
    matches = []
    for target in targets:
        if port not in target.get("ports", []):
            continue
        probe_results = []
        for probe in target.get("probes", []):
            result = run_probe(ip, port, probe)
            probe_results.append((probe, result))
        hit_probes = match_target(ip, port, probe_results)
        for hp in hit_probes:
            matches.append({
                "ip": ip,
                "port": port,
                "target_name": target["name"],
                "tags": target.get("tags", []),
                **hp,
            })
    return matches


# ── Commands ──────────────────────────────────────────────────────────────────
def cmd_run(args):
    banner()

    missing = check_deps()
    if missing:
        err(f"Missing required tools: {', '.join(missing)}")
        err("Install with: sudo apt install " + " ".join(missing))
        sys.exit(1)

    countries_path = Path(args.countries) if args.countries else DEFAULT_COUNTRIES
    targets_path = Path(args.targets) if args.targets else DEFAULT_TARGETS

    country_data = load_countries(countries_path)
    targets = load_targets(targets_path)

    profile = country_data.get("scan_profile", {})
    rate = profile.get("rate", 1000)
    max_hosts = profile.get("max_hosts_per_country", None)
    countries = sorted(
        country_data.get("countries", []),
        key=lambda c: {"high": 0, "medium": 1, "low": 2}.get(c.get("priority", "medium"), 1)
    )

    info(f"Engagement : {BOLD}{args.engagement}{RST}")
    info(f"Countries  : {', '.join(c['code'] for c in countries)}")
    info(f"Targets    : {len(targets)}")
    info(f"Rate       : {rate} pps")
    print()

    info("Loading country CIDR data...")
    fetch_rir_data()
    cidr_index = load_index()

    conn = db.get_db(args.engagement)
    run_id = db.start_run(
        conn,
        [c["code"] for c in countries],
        [t["name"] for t in targets],
    )

    all_ports = collect_all_ports(targets)
    info(f"Scanning ports: {', '.join(str(p) for p in all_ports)}")
    print()

    total_hosts = 0
    total_matches = 0
    workers = args.workers
    n_countries = len(countries)
    _interrupted = False

    def progress(msg: str, end="\n"):
        print(f"  {W}{msg}{RST}", end=end, flush=True)

    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)

            for idx, country in enumerate(countries, 1):
                cc = country["code"].upper()
                exclude = country.get("exclude_cidrs", [])
                cidrs = get_cidrs(cc, cidr_index)
                cidrs = merge_cidrs(cidrs, exclude)

                if not cidrs:
                    warn(f"{cc}: no CIDR ranges found, skipping")
                    continue

                if max_hosts:
                    cidrs = cidrs[:max_hosts // 256 + 1]

                total_ips = count_ips(cidrs)
                eta = fmt_eta(total_ips, len(all_ports), rate)
                print(f"\n{C}[{idx}/{n_countries}]{RST} {BOLD}{cc}{RST} — {len(cidrs)} CIDR blocks  {W}({fmt_ip_count(total_ips)} IPs  ~{eta} at {rate} pps){RST}")
                info(f"{cc}: running masscan  (Ctrl-C once = skip country, twice = abort all)...")

                try:
                    open_ports = masscan_sweep(cidrs, all_ports, rate, tmp)
                    _interrupted = False
                except KeyboardInterrupt:
                    print()
                    if _interrupted:
                        raise
                    _interrupted = True
                    warn(f"{cc}: skipping — press Ctrl-C again to abort entire scan")
                    continue

                hosts_found = len(set(e["ip"] for e in open_ports))
                total_hosts += hosts_found
                ok(f"{cc}: {hosts_found} live hosts, {len(open_ports)} open ports found")

                if not open_ports:
                    continue

                # Group by IP for nmap
                ip_ports: dict[str, list[int]] = {}
                for entry in open_ports:
                    ip_ports.setdefault(entry["ip"], []).append(entry["port"])

                total_nmap = len(ip_ports)
                info(f"{cc}: fingerprinting {total_nmap} hosts with nmap...")
                nmap_results: dict[str, dict] = {}
                nmap_done = 0
                with ThreadPoolExecutor(max_workers=min(workers, 10)) as ex:
                    futs = {ex.submit(nmap_fingerprint, ip, ports, tmp): ip
                            for ip, ports in ip_ports.items()}
                    for fut in as_completed(futs):
                        ip = futs[fut]
                        nmap_done += 1
                        progress(f"nmap: {nmap_done}/{total_nmap}", end="\r")
                        try:
                            nmap_results[ip] = fut.result()
                        except Exception:
                            nmap_results[ip] = {}
                print()

                total_probes = len(open_ports)
                info(f"{cc}: probing {total_probes} open ports for target matches...")
                probe_done = 0
                with ThreadPoolExecutor(max_workers=workers) as ex:
                    futs = [ex.submit(probe_host, entry["ip"], entry["port"], targets)
                            for entry in open_ports]

                    for fut in as_completed(futs):
                        probe_done += 1
                        progress(f"probe: {probe_done}/{total_probes}", end="\r")
                        try:
                            matches = fut.result()
                        except Exception:
                            continue
                        for match in matches:
                            ip = match["ip"]
                            port = match["port"]
                            nmap_info = nmap_results.get(ip, {}).get(port, {})
                            hit = {
                                **match,
                                "country_code": cc,
                                "nmap_service": nmap_info.get("service", ""),
                                "nmap_banner": nmap_info.get("banner", ""),
                            }
                            db.insert_hit(conn, run_id, hit)
                            total_matches += 1
                            ver = f" [{match['version']}]" if match.get("version") else ""
                            print(f"\r  {G}HIT{RST} {BOLD}{ip}:{port}{RST}  {Y}{match['target_name']}{RST}{ver}  [{match['confidence']}]  {cc}")
                print()

    except KeyboardInterrupt:
        print()
        warn("Scan aborted by user")

    db.finish_run(conn, run_id, total_hosts, total_matches)
    conn.close()

    print()
    ok(f"Run complete. {total_hosts} hosts scanned, {G}{BOLD}{total_matches}{RST} matches.")
    info(f"Results: results/{args.engagement}.db  (run_id={run_id})")


def cmd_export(args):
    conn = db.get_db(args.engagement)
    fmt = args.format.lower()

    if fmt == "json":
        data = db.export_json(conn)
        out = json.dumps(data, indent=2)
    elif fmt == "csv":
        lines = db.export_csv_lines(conn)
        out = "\n".join(lines)
    else:
        err(f"Unknown format: {fmt}")
        sys.exit(1)

    if args.output:
        Path(args.output).write_text(out)
        ok(f"Exported {fmt.upper()} → {args.output}")
    else:
        print(out)

    conn.close()


def cmd_update_cidr(args):
    banner()
    info("Force-refreshing RIR delegated data...")
    fetch_rir_data(force=True)
    from lib.cidr import load_index
    info("Rebuilding CIDR index...")
    index = load_index(force_rebuild=True)
    country_count = len(index)
    total_cidrs = sum(len(v) for v in index.values())
    ok(f"Index built: {country_count} countries, {total_cidrs} CIDR blocks")


def cmd_list(args):
    conn = db.get_db(args.engagement)
    rows = conn.execute(
        "SELECT target_name, country_code, COUNT(*) as c FROM hits GROUP BY target_name, country_code ORDER BY c DESC"
    ).fetchall()
    if not rows:
        warn(f"No hits in engagement: {args.engagement}")
        return
    print(f"\n{'Target':<40} {'Country':<8} {'Hits':>6}")
    print("-" * 58)
    for r in rows:
        print(f"  {r['target_name']:<38} {r['country_code']:<8} {r['c']:>6}")
    conn.close()


# ── Entry point ───────────────────────────────────────────────────────────────
def main():
    import argparse

    parser = argparse.ArgumentParser(
        prog="geo-scout",
        description="Geo-targeted network scanner"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # run
    p_run = sub.add_parser("run", help="Run a scan")
    p_run.add_argument("--engagement", "-e", required=True, help="Engagement name (used for DB filename)")
    p_run.add_argument("--countries", "-c", default=None, help="Path to countries.yaml")
    p_run.add_argument("--targets", "-t", default=None, help="Path to targets.yaml")
    p_run.add_argument("--workers", "-w", type=int, default=20, help="Probe thread workers (default: 20)")

    # export
    p_exp = sub.add_parser("export", help="Export results")
    p_exp.add_argument("--engagement", "-e", required=True)
    p_exp.add_argument("--format", "-f", default="csv", choices=["csv", "json"])
    p_exp.add_argument("--output", "-o", default=None, help="Output file (stdout if omitted)")

    # update-cidr
    sub.add_parser("update-cidr", help="Force refresh RIR CIDR data")

    # list
    p_list = sub.add_parser("list", help="Summarize hits for an engagement")
    p_list.add_argument("--engagement", "-e", required=True)

    args = parser.parse_args()

    if args.command == "run":
        cmd_run(args)
    elif args.command == "export":
        cmd_export(args)
    elif args.command == "update-cidr":
        cmd_update_cidr(args)
    elif args.command == "list":
        cmd_list(args)


if __name__ == "__main__":
    main()