#!/usr/bin/env python3
"""Supply chain infrastructure detection — passively identifies internal
package repos, update servers, config management, CI/CD, and container
registries by analyzing DNS and HTTP traffic patterns.

Detection only — no exploitation or active probing.
"""

import json
import logging
import os
import re
import sqlite3
import threading
import time
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS supply_chain (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    type            TEXT    NOT NULL,
    target_ip       TEXT    NOT NULL,
    target_hostname TEXT    DEFAULT '',
    target_port     INTEGER,
    evidence        TEXT    NOT NULL,
    confidence      REAL    DEFAULT 0.5,
    first_seen      REAL    NOT NULL,
    last_seen       REAL    NOT NULL,
    details         TEXT    DEFAULT '',
    UNIQUE(type, target_ip, target_port)
);

CREATE INDEX IF NOT EXISTS idx_sc_type ON supply_chain(type);
CREATE INDEX IF NOT EXISTS idx_sc_ip   ON supply_chain(target_ip);
"""

# -----------------------------------------------------------------------
# Detection signatures
# -----------------------------------------------------------------------

# HTTP path patterns -> (infra_type, description, confidence)
_HTTP_PATH_PATTERNS = [
    # Package repositories
    (re.compile(r"/simple/[\w-]+/"), "pypi_mirror", "PyPI mirror (PEP 503 simple API)", 0.9),
    (re.compile(r"/pypi/[\w-]+/json"), "pypi_mirror", "PyPI mirror (JSON API)", 0.9),
    (re.compile(r"/-/npm/v1/"), "npm_registry", "npm registry (v1 API)", 0.9),
    (re.compile(r"/@[\w-]+/[\w-]+/-/"), "npm_registry", "npm scoped package download", 0.85),
    (re.compile(r"/repository/maven-"), "maven_repo", "Maven/Nexus repository", 0.85),
    (re.compile(r"/artifactory/"), "artifactory", "JFrog Artifactory", 0.9),
    (re.compile(r"/content/repositories/"), "nexus_repo", "Sonatype Nexus repository", 0.85),
    (re.compile(r"/api/v4/projects/\d+/packages/"), "gitlab_packages", "GitLab Package Registry", 0.9),
    (re.compile(r"/v2/[\w.-]+/[\w.-]+/manifests/"), "container_registry", "Docker/OCI registry (manifest pull)", 0.95),
    (re.compile(r"/v2/[\w.-]+/[\w.-]+/blobs/"), "container_registry", "Docker/OCI registry (blob pull)", 0.95),
    (re.compile(r"/v2/_catalog"), "container_registry", "Docker registry catalog", 0.95),

    # WSUS
    (re.compile(r"/ClientWebService/"), "wsus", "WSUS client web service", 0.9),
    (re.compile(r"/SimpleAuthWebService/"), "wsus", "WSUS auth service", 0.9),
    (re.compile(r"/ApiRemoting30/"), "wsus", "WSUS API remoting", 0.9),
    (re.compile(r"/Content/[\w-]+\.cab"), "wsus", "WSUS content download", 0.8),

    # CI/CD
    (re.compile(r"/job/[\w-]+/build"), "jenkins", "Jenkins build trigger", 0.9),
    (re.compile(r"/api/json\?tree="), "jenkins", "Jenkins API query", 0.85),
    (re.compile(r"/queue/api/json"), "jenkins", "Jenkins queue API", 0.85),
    (re.compile(r"/api/v4/runners/"), "gitlab_ci", "GitLab CI runner API", 0.9),
    (re.compile(r"/api/v3/repos/.+/actions/"), "github_actions", "GitHub Actions API", 0.85),
    (re.compile(r"/_apis/build/"), "azure_devops", "Azure DevOps build API", 0.9),

    # Config management
    (re.compile(r"/puppet/v3/"), "puppet", "Puppet v3 API", 0.9),
    (re.compile(r"/puppet/v4/"), "puppet", "Puppet v4 API", 0.9),
    (re.compile(r"/organizations/[\w-]+/cookbooks/"), "chef", "Chef cookbook API", 0.9),
    (re.compile(r"/nodes/[\w.-]+"), "chef", "Chef node API", 0.8),

    # Internal CA
    (re.compile(r"/certsrv/"), "internal_ca", "Microsoft AD CS certificate services", 0.9),
    (re.compile(r"/acme/directory"), "internal_ca", "ACME CA directory (internal Let's Encrypt/step-ca)", 0.8),
    (re.compile(r"/v1/pki/"), "vault_pki", "HashiCorp Vault PKI engine", 0.9),

    # SCCM / MECM
    (re.compile(r"/SMS_MP/"), "sccm", "SCCM Management Point", 0.9),
    (re.compile(r"/CCM_Client/"), "sccm", "SCCM client endpoint", 0.9),
    (re.compile(r"/CMGateway/"), "sccm", "SCCM Cloud Management Gateway", 0.85),
]

# DNS patterns -> (infra_type, description, confidence)
_DNS_PATTERNS = [
    (re.compile(r"(pypi|pip)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "pypi_mirror", "Internal PyPI mirror hostname", 0.85),
    (re.compile(r"(npm|registry)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "npm_registry", "Internal npm registry hostname", 0.85),
    (re.compile(r"(nexus|artifactory|repo)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "artifact_repo", "Internal artifact repository", 0.8),
    (re.compile(r"(docker|registry|harbor|quay)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "container_registry", "Internal container registry", 0.85),
    (re.compile(r"(wsus|update)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "wsus", "WSUS/update server", 0.8),
    (re.compile(r"(jenkins|ci|build|drone|bamboo|teamcity|concourse)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "ci_cd", "CI/CD server", 0.8),
    (re.compile(r"(gitlab)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "gitlab", "GitLab instance", 0.9),
    (re.compile(r"(puppet|puppetdb|foreman)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "puppet", "Puppet infrastructure", 0.85),
    (re.compile(r"(chef|supermarket)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "chef", "Chef infrastructure", 0.85),
    (re.compile(r"(ansible|tower|awx)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "ansible", "Ansible Tower/AWX", 0.85),
    (re.compile(r"(sccm|mecm|configmgr)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "sccm", "SCCM/MECM infrastructure", 0.9),
    (re.compile(r"(ca|pki|cert|issuing)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "internal_ca", "Internal CA/PKI", 0.7),
    (re.compile(r"(vault)[\w-]*\.[\w.]+\.(local|internal|corp|lan)$", re.I),
     "hashicorp_vault", "HashiCorp Vault", 0.8),
]

# Port patterns -> (infra_type, description, confidence)
_PORT_PATTERNS = {
    8140: ("puppet", "Puppet master default port", 0.7),
    8443: ("ci_cd", "Common CI/CD HTTPS port", 0.3),  # low confidence — many things use 8443
    8080: ("jenkins", "Jenkins default HTTP port", 0.3),
    50000: ("jenkins", "Jenkins agent JNLP port", 0.7),
    5000: ("container_registry", "Docker registry default port", 0.5),
    443: ("generic_https", "HTTPS", 0.1),  # too common for high confidence alone
    8140: ("puppet", "Puppet master", 0.7),
    4433: ("chef", "Chef Automate", 0.6),
}


class SupplyChainDetect(BaseModule):
    """Passively detect internal supply chain infrastructure."""

    name = "supply_chain_detect"
    module_type = "intel"
    priority = 200
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._db_path = ""
        self._conn: Optional[sqlite3.Connection] = None
        self._lock = threading.Lock()
        self._scan_thread: Optional[threading.Thread] = None
        self._detection_count = 0

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        base_dir = self.config.get("data_dir", os.path.expanduser("~/.implant"))
        self._db_path = os.path.join(base_dir, "supply_chain.db")
        os.makedirs(os.path.dirname(self._db_path), exist_ok=True)

        self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(_SCHEMA)

        # Subscribe to host discovery for port-based detection
        self.bus.subscribe(self._on_host_discovered, "HOST_DISCOVERED")

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()

        # Periodic scan of accumulated DNS and HTTP data
        self._scan_thread = threading.Thread(
            target=self._periodic_scan, daemon=True,
            name="sensor-supply-chain-scan",
        )
        self._scan_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("SupplyChainDetect started — db=%s", self._db_path)

    def stop(self) -> None:
        if not self._running:
            return
        self._running = False

        self.bus.unsubscribe(self._on_host_discovered, "HOST_DISCOVERED")

        if self._conn:
            self._conn.close()
            self._conn = None

        self.state.set_module_status(self.name, "stopped")
        logger.info("SupplyChainDetect stopped — %d detections", self._detection_count)

    def status(self) -> dict:
        det_count = 0
        if self._conn:
            try:
                with self._lock:
                    row = self._conn.execute(
                        "SELECT COUNT(*) FROM supply_chain"
                    ).fetchone()
                    det_count = row[0] if row else 0
            except Exception:
                pass

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "detections": det_count,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # Detection functions
    # ------------------------------------------------------------------

    def analyze_http_request(self, target_ip: str, target_port: int,
                             method: str, path: str, host_header: str = "",
                             user_agent: str = "") -> Optional[dict]:
        """Analyze an HTTP request for supply chain infrastructure indicators.

        Called by tool_output_parser or directly by HTTP monitoring modules.

        Returns:
            Detection dict if found, None otherwise.
        """
        for pattern, infra_type, description, confidence in _HTTP_PATH_PATTERNS:
            if pattern.search(path):
                det = self._record_detection(
                    infra_type=infra_type,
                    target_ip=target_ip,
                    target_hostname=host_header,
                    target_port=target_port,
                    evidence=f"{method} {path}",
                    confidence=confidence,
                    details=json.dumps({
                        "method": method,
                        "path": path,
                        "host": host_header,
                        "user_agent": user_agent,
                        "description": description,
                    }),
                )
                return det
        return None

    def analyze_dns_query(self, query_name: str, resolved_ip: str = "",
                          client_ip: str = "") -> Optional[dict]:
        """Analyze a DNS query for supply chain infrastructure indicators.

        Returns:
            Detection dict if found, None otherwise.
        """
        for pattern, infra_type, description, confidence in _DNS_PATTERNS:
            if pattern.search(query_name):
                det = self._record_detection(
                    infra_type=infra_type,
                    target_ip=resolved_ip or "unresolved",
                    target_hostname=query_name,
                    target_port=None,
                    evidence=f"DNS: {query_name}",
                    confidence=confidence,
                    details=json.dumps({
                        "query": query_name,
                        "resolved_ip": resolved_ip,
                        "client_ip": client_ip,
                        "description": description,
                    }),
                )
                return det
        return None

    def analyze_port(self, target_ip: str, port: int,
                     hostname: str = "") -> Optional[dict]:
        """Check if a discovered open port suggests supply chain infrastructure."""
        if port not in _PORT_PATTERNS:
            return None

        infra_type, description, confidence = _PORT_PATTERNS[port]
        if confidence < 0.5:
            # Low-confidence port-only detections need corroboration
            # Check if we already have DNS or HTTP evidence for this IP
            existing = self._get_detections_for_ip(target_ip)
            if existing:
                # Boost confidence if corroborated
                confidence = min(confidence + 0.3, 0.95)
            else:
                return None

        return self._record_detection(
            infra_type=infra_type,
            target_ip=target_ip,
            target_hostname=hostname,
            target_port=port,
            evidence=f"Open port {port}",
            confidence=confidence,
            details=json.dumps({"description": description, "port": port}),
        )

    def _on_host_discovered(self, event) -> None:
        """Check discovered host ports for supply chain indicators."""
        p = event.payload
        ip = p.get("ip", "")
        hostname = p.get("hostname", "")
        ports = p.get("open_ports", [])

        for port in ports:
            self.analyze_port(ip, port, hostname)

    # ------------------------------------------------------------------
    # Database operations
    # ------------------------------------------------------------------

    def _record_detection(self, infra_type: str, target_ip: str,
                          target_hostname: str, target_port: Optional[int],
                          evidence: str, confidence: float,
                          details: str = "") -> dict:
        """Record a supply chain detection, deduplicating by type+ip+port."""
        now = time.time()
        detection = {
            "type": infra_type,
            "target_ip": target_ip,
            "target_hostname": target_hostname,
            "target_port": target_port,
            "evidence": evidence,
            "confidence": confidence,
            "details": details,
        }

        with self._lock:
            try:
                self._conn.execute(
                    """INSERT INTO supply_chain
                       (type, target_ip, target_hostname, target_port,
                        evidence, confidence, first_seen, last_seen, details)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                       ON CONFLICT(type, target_ip, target_port) DO UPDATE SET
                           last_seen = excluded.last_seen,
                           confidence = MAX(supply_chain.confidence, excluded.confidence),
                           evidence = supply_chain.evidence || '; ' || excluded.evidence,
                           target_hostname = CASE WHEN excluded.target_hostname != ''
                                              THEN excluded.target_hostname
                                              ELSE supply_chain.target_hostname END
                    """,
                    (infra_type, target_ip, target_hostname, target_port,
                     evidence, confidence, now, now, details),
                )
                self._conn.commit()
                self._detection_count += 1

                logger.info(
                    "Supply chain detection: %s at %s:%s (%s, confidence=%.2f)",
                    infra_type, target_ip, target_port, evidence, confidence,
                )
            except Exception:
                logger.exception("Failed to record supply chain detection")

        return detection

    def _get_detections_for_ip(self, ip: str) -> list:
        """Get all detections for a given IP."""
        with self._lock:
            try:
                rows = self._conn.execute(
                    "SELECT * FROM supply_chain WHERE target_ip = ?", (ip,)
                ).fetchall()
                return [dict(r) for r in rows]
            except Exception:
                return []

    # ------------------------------------------------------------------
    # Periodic scanning
    # ------------------------------------------------------------------

    def _periodic_scan(self) -> None:
        """Periodically pull DNS and HTTP data from other modules' state."""
        interval = self.config.get("supply_chain_scan_interval", 180)
        while self._running:
            time.sleep(interval)
            try:
                self._scan_dns_data()
                self._scan_http_data()
            except Exception:
                logger.exception("Periodic supply chain scan failed")

    def _scan_dns_data(self) -> None:
        """Pull recent DNS queries from dns_logger state and analyze."""
        last_ts = self.state.get(self.name, "dns_scan_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        dns_json = self.state.get("dns_logger", "recent_queries")
        if not dns_json:
            return

        try:
            queries = json.loads(dns_json)
            max_ts = last_ts
            for q in queries:
                ts = q.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.analyze_dns_query(
                    query_name=q.get("query_name", ""),
                    resolved_ip=q.get("resolved_ip", ""),
                    client_ip=q.get("client_ip", ""),
                )
                max_ts = max(max_ts, ts)

            self.state.set(self.name, "dns_scan_last_ts", str(max_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    def _scan_http_data(self) -> None:
        """Pull recent HTTP log data from state and analyze."""
        last_ts = self.state.get(self.name, "http_scan_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        http_json = self.state.get("http_monitor", "recent_requests")
        if not http_json:
            return

        try:
            requests = json.loads(http_json)
            max_ts = last_ts
            for req in requests:
                ts = req.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.analyze_http_request(
                    target_ip=req.get("target_ip", ""),
                    target_port=req.get("target_port", 80),
                    method=req.get("method", "GET"),
                    path=req.get("path", ""),
                    host_header=req.get("host", ""),
                    user_agent=req.get("user_agent", ""),
                )
                max_ts = max(max_ts, ts)

            self.state.set(self.name, "http_scan_last_ts", str(max_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    # ------------------------------------------------------------------
    # Query interface
    # ------------------------------------------------------------------

    def get_detections(self, infra_type: str = None,
                       min_confidence: float = 0.0) -> list:
        """Get supply chain detections, optionally filtered."""
        with self._lock:
            if infra_type:
                rows = self._conn.execute(
                    """SELECT * FROM supply_chain
                       WHERE type = ? AND confidence >= ?
                       ORDER BY confidence DESC""",
                    (infra_type, min_confidence),
                ).fetchall()
            else:
                rows = self._conn.execute(
                    """SELECT * FROM supply_chain
                       WHERE confidence >= ?
                       ORDER BY confidence DESC""",
                    (min_confidence,),
                ).fetchall()

        return [dict(r) for r in rows]

    def get_summary(self) -> dict:
        """Return a summary of detected supply chain infrastructure."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT type, COUNT(*) as count,
                          AVG(confidence) as avg_confidence
                   FROM supply_chain
                   GROUP BY type
                   ORDER BY count DESC"""
            ).fetchall()

        return {
            "total_detections": sum(r[1] for r in rows),
            "by_type": {r[0]: {"count": r[1], "avg_confidence": round(r[2], 2)}
                        for r in rows},
        }
