main.py

#!/usr/bin/env python3
"""
CyberRemedy v1.2 — Entry point

Usage:
  python main.py                    # Start API + dashboard (default)
  python main.py --port 9000        # Custom port
  python main.py --host 127.0.0.1   # Bind to localhost only
  python main.py --train            # Train ML models
"""
import argparse, logging, logging.handlers, sys, os
from pathlib import Path

ROOT = Path(__file__).parent

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(name)-26s %(levelname)-8s %(message)s",
)
logger = logging.getLogger("cyberremedy.main")


def ensure_dirs():
    for d in [
        "data", "data/reports", "data/datasets",
        "data/yara_rules", "data/sigma_rules",
        "data/lake", "data/lake/hot", "data/lake/warm", "data/lake/cold",
        "data/logs", "data/logs/alerts", "data/logs/traffic",
        "data/logs/blocks", "data/logs/events", "data/logs/assets",
        "data/pcap", "data/geoip", "data/assets",
        "models", "logs",
    ]:
        (ROOT / d).mkdir(parents=True, exist_ok=True)


def setup_file_logging():
    """Rotate log at 10 MB, keep 10 files. All DEBUG+ goes to file."""
    log_path = ROOT / "data" / "logs" / "cyberremedy.log"
    fh = logging.handlers.RotatingFileHandler(
        log_path, maxBytes=10 * 1024 * 1024, backupCount=10, encoding="utf-8"
    )
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter(
        "%(asctime)s %(name)-26s %(levelname)-8s %(message)s"
    ))
    logging.getLogger().addHandler(fh)
    logging.getLogger().setLevel(logging.DEBUG)
    logger.info(f"File logging active → {log_path}")


def download_geoip():
    """
    Download free offline GeoIP data from db-ip.com — no API key needed.
    Only runs if the file doesn't already exist.
    """
    import urllib.request, gzip as _gz, csv, json
    from datetime import datetime as _dt

    out_json = ROOT / "data" / "geoip" / "ip_country.json"
    marker   = ROOT / "data" / "geoip" / "ip_country.csv"

    import time as _geoip_time
    _geoip_age = _geoip_time.time() - marker.stat().st_mtime if marker.exists() else 999999
    if marker.exists() and marker.stat().st_size > 50_000 and _geoip_age < 30 * 86400:
        logger.info(f"GeoIP: offline DB present ({marker.stat().st_size//1024} KB, {int(_geoip_age/86400)}d old) — skipping download")
        return
    if marker.exists() and _geoip_age >= 30 * 86400:
        logger.info("GeoIP: database is 30+ days old — auto-refreshing…")

    # Try multiple free sources (no account or key needed)
    sources = [
        # Primary: ip-location-db GeoIP CSV (GitHub raw, no key, updated monthly)
        "https://raw.githubusercontent.com/sapics/ip-location-db/main/geo-whois-asn-country/geo-whois-asn-country-ipv4.csv",
        # Fallback 1: country-only CSV from same repo
        "https://raw.githubusercontent.com/sapics/ip-location-db/main/asn-country/asn-country-ipv4.csv",
        # Fallback 2: dbip-country via jsdelivr CDN (very reliable)
        "https://cdn.jsdelivr.net/npm/@ip-location-db/geo-whois-asn-country/geo-whois-asn-country-ipv4.csv",
        # Fallback 3: country ipv4 via jsDelivr
        "https://cdn.jsdelivr.net/npm/@ip-location-db/asn-country/asn-country-ipv4.csv",
    ]

    for url in sources:
        try:
            logger.info(f"GeoIP: trying {url}")
            import zipfile, io
            with urllib.request.urlopen(url, timeout=30) as resp:
                raw = resp.read()
            # Handle zip vs csv
            if url.endswith(".ZIP") or url.endswith(".zip"):
                zf = zipfile.ZipFile(io.BytesIO(raw))
                data = zf.read(zf.namelist()[0]).decode("utf-8", errors="replace")
            elif url.endswith(".gz"):
                data = _gz.decompress(raw).decode("utf-8", errors="replace")
            else:
                data = raw.decode("utf-8", errors="replace")
            raw_rows = list(csv.reader(data.splitlines()))
            # Handle both formats: (start,end,cc) and (start_num,end_num,cc)
            rows = [r[:3] for r in raw_rows if len(r) >= 3 and not r[0].startswith('#')]
            json.dump(rows, open(out_json, "w"))
            marker.write_text(f"source:{url}\nrows:{len(rows)}")
            logger.info(f"GeoIP: saved {len(rows):,} IP ranges → {out_json}")
            return
        except Exception as e:
            logger.warning(f"GeoIP source failed ({url}): {e}")

    logger.warning("GeoIP: offline download failed — ip-api.com online fallback will be used (45 req/min free)")


def _ensure_packages():
    """Auto-install missing critical packages before starting."""
    import subprocess, importlib
    REQUIRED = [
        ("fastapi",   "fastapi==0.111.0"),
        ("uvicorn",   "uvicorn[standard]==0.29.0"),
        ("pydantic",  "pydantic>=2.0"),
        ("aiofiles",  "aiofiles>=23.1"),
        ("websockets","websockets>=12.0"),
    ]
    missing = []
    for mod, pkg in REQUIRED:
        try:
            importlib.import_module(mod)
        except ImportError:
            missing.append(pkg)
    if not missing:
        return
    print(f"\n⚙  Auto-installing {len(missing)} missing package(s): {', '.join(missing)}")
    print("   (Run:  sudo pip3 install -r requirements.txt  to install everything at once)\n")
    for pkg in missing:
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", pkg,
                 "--quiet", "--break-system-packages"],
                stderr=subprocess.DEVNULL
            )
            print(f"   ✅ {pkg}")
        except subprocess.CalledProcessError:
            try:
                subprocess.check_call(
                    [sys.executable, "-m", "pip", "install", pkg, "--quiet"],
                    stderr=subprocess.DEVNULL
                )
                print(f"   ✅ {pkg}")
            except Exception:
                print(f"   ❌ Failed to install {pkg}")
                print(f"      Fix: sudo pip3 install {pkg}")
    # Re-check
    still_missing = []
    for mod, pkg in REQUIRED:
        try:
            importlib.import_module(mod)
        except ImportError:
            still_missing.append(pkg)
    if still_missing:
        print(f"\n❌ Still missing after install: {still_missing}")
        print("   Run:  sudo pip3 install -r requirements.txt")
        print("   Or:   sudo bash install_and_run.sh")
        sys.exit(1)
    print("   ✅ All packages ready\n")

def cmd_api(args):
    ensure_dirs()
    setup_file_logging()
    _ensure_packages()
    try:
        download_geoip()
    except Exception as e:
        logger.warning(f"GeoIP setup skipped: {e}")

    import uvicorn
    logger.info(f"Starting CyberRemedy v1.2 on http://{args.host}:{args.port}")
    print(f"\n  ✅ CyberRemedy v1.2 running")
    print(f"  Dashboard → http://{args.host if args.host != '0.0.0.0' else 'localhost'}:{args.port}")
    print(f"  API Docs  → http://{args.host if args.host != '0.0.0.0' else 'localhost'}:{args.port}/docs\n")
    uvicorn.run(
        "api.server:app",
        host=args.host,
        port=args.port,
        reload=False,
        log_level="warning",
        app_dir=str(ROOT),
    )


def cmd_train(args):
    ensure_dirs()
    logger.info("Training ML models …")
    sys.path.insert(0, str(ROOT))

    # ── Auto-install sklearn/joblib if missing ─────────────────────────────
    try:
        import sklearn  # noqa: F401
    except ImportError:
        logger.info("scikit-learn not found — installing automatically …")
        import subprocess
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", "scikit-learn", "joblib",
             "--break-system-packages", "--quiet"],
            capture_output=True, text=True
        )
        if result.returncode != 0:
            logger.warning(f"pip install failed: {result.stderr[:300]}")
        else:
            logger.info("scikit-learn installed successfully")

    from detection.anomaly import AnomalyDetector
    import numpy as np

    detector = AnomalyDetector(
        model_path="models/anomaly_model.joblib",
        classifier_path="models/rf_attack_model.joblib",
    )
    FEATURES = [
        "packet_count","total_bytes","bytes_per_second","packets_per_second",
        "avg_packet_size","min_packet_size","max_packet_size","std_packet_size",
        "flow_duration","avg_inter_arrival","std_inter_arrival","min_inter_arrival",
        "unique_dst_ports","unique_src_ips","dst_port_entropy","flag_entropy",
        "ttl_entropy","payload_entropy","has_syn","has_fin","has_rst","has_null",
    ]
    rng = np.random.default_rng(42)
    benign = [{f: float(rng.uniform(0,1)) for f in FEATURES} for _ in range(500)]
    attack = [{f: float(rng.uniform(0,1)) for f in FEATURES} for _ in range(100)]
    for s in attack:
        s["packet_count"] = float(rng.integers(500, 5000))
        s["unique_dst_ports"] = float(rng.integers(50, 200))
    all_flows = benign + attack
    labels = [0]*500 + [1]*100
    detector.train(all_flows, labels)
    logger.info("Isolation Forest + RandomForest saved to models/")

    # ── Train LSTM if not already trained ─────────────────────────────────
    from pathlib import Path as _P
    lstm_trained = _P("models/lstm_detector.pt").exists() or _P("models/lstm_detector.keras").exists()
    force_lstm = getattr(args, "_force_lstm", False)
    if not lstm_trained or force_lstm:
        logger.info("Training LSTM (NSL-KDD) — this takes ~30 s on CPU …")
        import subprocess
        extra = ["--retrain"] if force_lstm else []
        result = subprocess.run(
            [sys.executable, "ml/lstm/trainer.py"] + extra,
            capture_output=False, text=True, timeout=600
        )
        if result.returncode == 0:
            logger.info("LSTM training complete")
        else:
            logger.warning("LSTM training exited with non-zero code — check output above")
    else:
        logger.info("LSTM model already exists — skipping (use --retrain to force)")

    logger.info("Training complete — models saved to models/")



def main():
    p = argparse.ArgumentParser(description="CyberRemedy v1.2 SIEM")
    p.add_argument("--host",    default="0.0.0.0")
    p.add_argument("--port",    type=int, default=8000)
    p.add_argument("--train",   action="store_true", help="Train ML models and exit")
    p.add_argument("--retrain", action="store_true", help="Force retrain all ML models (overwrite existing)")
    args = p.parse_args()

    if args.retrain:
        args.train = True   # retrain implies train
        args._force_lstm = True
    else:
        args._force_lstm = False

    if args.train:
        cmd_train(args)
    else:
        cmd_api(args)


if __name__ == "__main__":
    main()