"""
MundiX Knowledge Graph
=======================

A graph-based data structure that tracks relationships between discovered
assets during a pentesting engagement:

    Target → Host → Port → Service → Version → CVE → Exploit → Credential → Shell

The KG enables:
  1. Attack chain discovery (shortest path from external to root shell)
  2. AI context enrichment (feed the graph state to the LLM)
  3. Priority scoring (which paths are most exploitable?)
  4. Report generation (visual attack narrative)

Usage:
    from knowledge_graph import KnowledgeGraph, NodeType
    
    kg = KnowledgeGraph("example.com")
    kg.add_host("10.0.0.1", source="nmap")
    kg.add_port("10.0.0.1", 80, "tcp", "http", "Apache/2.4.41")
    kg.add_cve("10.0.0.1", 80, "CVE-2021-41773", cvss=9.8, severity="critical")
    kg.add_exploit("CVE-2021-41773", "EDB-50383", exploit_type="path_traversal")
    
    chains = kg.find_attack_chains()
    context = kg.to_ai_context()
"""

import json
import os
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Dict, List, Optional, Set, Tuple, Any


# ---------------------------------------------------------------------------
# Node Types & Edge Types
# ---------------------------------------------------------------------------

class NodeType(Enum):
    TARGET = "target"
    HOST = "host"
    PORT = "port"
    SERVICE = "service"
    CVE = "cve"
    EXPLOIT = "exploit"
    CREDENTIAL = "credential"
    SHELL = "shell"
    DOMAIN = "domain"
    SUBDOMAIN = "subdomain"
    TECHNOLOGY = "technology"
    WAF = "waf"
    EMAIL = "email"


class EdgeType(Enum):
    RESOLVES_TO = "resolves_to"      # domain → host
    HAS_PORT = "has_port"            # host → port
    RUNS_SERVICE = "runs_service"    # port → service
    HAS_CVE = "has_cve"             # service → cve
    HAS_EXPLOIT = "has_exploit"      # cve → exploit
    YIELDS_CRED = "yields_cred"     # exploit → credential
    GIVES_SHELL = "gives_shell"     # credential/exploit → shell
    SUBDOMAIN_OF = "subdomain_of"   # subdomain → domain
    USES_TECH = "uses_tech"         # host → technology
    PROTECTED_BY = "protected_by"   # host → waf
    ASSOCIATED = "associated"        # email → domain
    LATERAL_TO = "lateral_to"       # shell → host (lateral movement)


# ---------------------------------------------------------------------------
# Graph Node
# ---------------------------------------------------------------------------

@dataclass
class Node:
    """A node in the knowledge graph."""
    id: str
    type: NodeType
    label: str
    properties: Dict[str, Any] = field(default_factory=dict)
    discovered_at: str = field(default_factory=lambda: datetime.now().isoformat())
    source: str = ""  # Which skill/tool discovered this
    confidence: float = 1.0  # 0-1, how confident we are

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other):
        return isinstance(other, Node) and self.id == other.id


@dataclass
class Edge:
    """A directed edge in the knowledge graph."""
    source_id: str
    target_id: str
    type: EdgeType
    properties: Dict[str, Any] = field(default_factory=dict)
    weight: float = 1.0  # Lower = more exploitable


# ---------------------------------------------------------------------------
# Knowledge Graph
# ---------------------------------------------------------------------------

class KnowledgeGraph:
    """
    A directed graph tracking pentest asset relationships.
    Designed for fast querying of attack chains and AI context generation.
    """

    def __init__(self, target: str):
        self.target = target
        self.nodes: Dict[str, Node] = {}
        self.edges: List[Edge] = []
        self._adjacency: Dict[str, List[Edge]] = defaultdict(list)  # Forward edges
        self._reverse: Dict[str, List[Edge]] = defaultdict(list)    # Reverse edges
        self.created_at = datetime.now().isoformat()

        # Create the root target node
        self._add_node(Node(
            id=f"target:{target}",
            type=NodeType.TARGET,
            label=target,
            properties={"original_target": target},
        ))

    # --- Node Addition Methods ---

    def _add_node(self, node: Node) -> Node:
        """Add a node (or update if exists)."""
        if node.id in self.nodes:
            # Merge properties
            self.nodes[node.id].properties.update(node.properties)
            return self.nodes[node.id]
        self.nodes[node.id] = node
        return node

    def _add_edge(self, source_id: str, target_id: str, edge_type: EdgeType,
                  weight: float = 1.0, **properties) -> Optional[Edge]:
        """Add an edge between two nodes."""
        if source_id not in self.nodes or target_id not in self.nodes:
            return None
        # Avoid duplicates
        for existing in self._adjacency[source_id]:
            if existing.target_id == target_id and existing.type == edge_type:
                existing.properties.update(properties)
                return existing
        edge = Edge(source_id, target_id, edge_type, properties, weight)
        self.edges.append(edge)
        self._adjacency[source_id].append(edge)
        self._reverse[target_id].append(edge)
        return edge

    def add_host(self, ip: str, hostname: str = None,
                 source: str = "nmap") -> str:
        """Add a discovered host."""
        node_id = f"host:{ip}"
        self._add_node(Node(
            id=node_id, type=NodeType.HOST, label=ip,
            properties={"ip": ip, "hostname": hostname or ""},
            source=source,
        ))
        self._add_edge(f"target:{self.target}", node_id, EdgeType.RESOLVES_TO)
        return node_id

    def add_domain(self, domain: str, source: str = "dns") -> str:
        """Add a discovered domain/subdomain."""
        node_id = f"domain:{domain}"
        self._add_node(Node(
            id=node_id, type=NodeType.SUBDOMAIN, label=domain,
            properties={"domain": domain},
            source=source,
        ))
        self._add_edge(f"target:{self.target}", node_id, EdgeType.SUBDOMAIN_OF)
        return node_id

    def add_port(self, host: str, port: int, protocol: str = "tcp",
                 service: str = "", version: str = "",
                 source: str = "nmap") -> str:
        """Add a discovered port with its service."""
        host_id = f"host:{host}"
        if host_id not in self.nodes:
            self.add_host(host, source=source)

        port_id = f"port:{host}:{port}/{protocol}"
        self._add_node(Node(
            id=port_id, type=NodeType.PORT, label=f"{port}/{protocol}",
            properties={"port": port, "protocol": protocol, "host": host},
            source=source,
        ))
        self._add_edge(host_id, port_id, EdgeType.HAS_PORT)

        # Add service node if known
        if service:
            svc_id = f"service:{host}:{port}:{service}"
            self._add_node(Node(
                id=svc_id, type=NodeType.SERVICE,
                label=f"{service} {version}".strip(),
                properties={"service": service, "version": version,
                            "host": host, "port": port},
                source=source,
            ))
            self._add_edge(port_id, svc_id, EdgeType.RUNS_SERVICE)

        return port_id

    def add_cve(self, host: str, port: int, cve_id: str,
                cvss: float = 0.0, severity: str = "unknown",
                description: str = "", source: str = "nvd") -> str:
        """Add a CVE linked to a service on a host:port."""
        node_id = f"cve:{cve_id}"
        self._add_node(Node(
            id=node_id, type=NodeType.CVE, label=cve_id,
            properties={"cvss": cvss, "severity": severity,
                        "description": description[:200]},
            source=source,
        ))

        # Link to service (find the service node for this host:port)
        svc_nodes = [n for n_id, n in self.nodes.items()
                     if n.type == NodeType.SERVICE
                     and n.properties.get("host") == host
                     and n.properties.get("port") == port]
        if svc_nodes:
            # Weight: higher CVSS = lower weight (more attractive path)
            weight = max(0.1, (10.0 - cvss) / 10.0)
            self._add_edge(svc_nodes[0].id, node_id, EdgeType.HAS_CVE, weight=weight)
        else:
            # Link directly to port
            port_id = f"port:{host}:{port}/tcp"
            if port_id in self.nodes:
                self._add_edge(port_id, node_id, EdgeType.HAS_CVE)

        return node_id

    def add_exploit(self, cve_id: str, exploit_id: str,
                    exploit_type: str = "unknown", title: str = "",
                    source: str = "exploitdb") -> str:
        """Add an exploit linked to a CVE."""
        node_id = f"exploit:{exploit_id}"
        self._add_node(Node(
            id=node_id, type=NodeType.EXPLOIT, label=title or exploit_id,
            properties={"exploit_type": exploit_type, "title": title,
                        "cve": cve_id},
            source=source,
        ))
        cve_node_id = f"cve:{cve_id}"
        if cve_node_id in self.nodes:
            self._add_edge(cve_node_id, node_id, EdgeType.HAS_EXPLOIT, weight=0.3)
        return node_id

    def add_credential(self, service: str, host: str, username: str,
                       password: str = "", hash_value: str = "",
                       source: str = "bruteforce") -> str:
        """Add a discovered credential."""
        node_id = f"cred:{host}:{service}:{username}"
        self._add_node(Node(
            id=node_id, type=NodeType.CREDENTIAL,
            label=f"{username}@{host}:{service}",
            properties={"username": username, "password": "***",
                        "hash": hash_value[:10] + "..." if hash_value else "",
                        "service": service, "host": host},
            source=source,
        ))
        # Link to the exploit or service that yielded it
        svc_nodes = [n for n in self.nodes.values()
                     if n.type == NodeType.SERVICE
                     and n.properties.get("host") == host]
        for svc in svc_nodes:
            self._add_edge(svc.id, node_id, EdgeType.YIELDS_CRED, weight=0.2)
        return node_id

    def add_shell(self, host: str, shell_type: str = "reverse",
                  privilege: str = "user", source: str = "") -> str:
        """Add a obtained shell/access."""
        node_id = f"shell:{host}:{privilege}"
        self._add_node(Node(
            id=node_id, type=NodeType.SHELL,
            label=f"{privilege} shell on {host}",
            properties={"shell_type": shell_type, "privilege": privilege,
                        "host": host},
            source=source,
        ))
        # Link from credentials on this host
        for n_id, n in self.nodes.items():
            if n.type == NodeType.CREDENTIAL and n.properties.get("host") == host:
                self._add_edge(n_id, node_id, EdgeType.GIVES_SHELL, weight=0.1)
            if n.type == NodeType.EXPLOIT:
                # Check if exploit's CVE is linked to this host
                for edge in self._adjacency.get(n_id, []):
                    if edge.target_id == node_id:
                        break
                else:
                    # Check if any CVE→exploit chain leads to this host
                    pass
        return node_id

    def add_technology(self, host: str, tech: str, version: str = "",
                       source: str = "whatweb") -> str:
        """Add a detected technology (CMS, framework, etc)."""
        node_id = f"tech:{host}:{tech}"
        self._add_node(Node(
            id=node_id, type=NodeType.TECHNOLOGY,
            label=f"{tech} {version}".strip(),
            properties={"tech": tech, "version": version},
            source=source,
        ))
        host_id = f"host:{host}"
        if host_id in self.nodes:
            self._add_edge(host_id, node_id, EdgeType.USES_TECH)
        return node_id

    def add_waf(self, host: str, waf_name: str, source: str = "wafw00f") -> str:
        """Add a detected WAF."""
        node_id = f"waf:{host}:{waf_name}"
        self._add_node(Node(
            id=node_id, type=NodeType.WAF, label=waf_name,
            properties={"waf": waf_name},
            source=source,
        ))
        host_id = f"host:{host}"
        if host_id in self.nodes:
            self._add_edge(host_id, node_id, EdgeType.PROTECTED_BY)
        return node_id

    # --- Query Methods ---

    def find_attack_chains(self, max_depth: int = 8) -> List[List[str]]:
        """
        Find all paths from target to shell/credential nodes.
        Returns attack chains sorted by total weight (lower = easier to exploit).
        """
        chains = []
        target_id = f"target:{self.target}"

        # Goal nodes: shells and credentials
        goal_types = {NodeType.SHELL, NodeType.CREDENTIAL, NodeType.EXPLOIT}
        goals = {n_id for n_id, n in self.nodes.items() if n.type in goal_types}

        if not goals:
            return []

        # BFS/DFS to find all paths
        def dfs(current: str, path: List[str], weight: float, visited: Set[str]):
            if len(path) > max_depth:
                return
            if current in goals and len(path) > 1:
                chains.append((list(path), weight))
                # Don't return — continue exploring for longer chains
            for edge in self._adjacency.get(current, []):
                if edge.target_id not in visited:
                    visited.add(edge.target_id)
                    path.append(edge.target_id)
                    dfs(edge.target_id, path, weight + edge.weight, visited)
                    path.pop()
                    visited.discard(edge.target_id)

        visited = {target_id}
        dfs(target_id, [target_id], 0.0, visited)

        # Sort by weight (lower = more exploitable) then by length (shorter preferred)
        chains.sort(key=lambda x: (x[1], len(x[0])))
        return [chain for chain, _ in chains[:20]]  # Top 20 chains

    def get_attack_surface_score(self) -> Dict[str, Any]:
        """Calculate attack surface metrics."""
        hosts = [n for n in self.nodes.values() if n.type == NodeType.HOST]
        ports = [n for n in self.nodes.values() if n.type == NodeType.PORT]
        services = [n for n in self.nodes.values() if n.type == NodeType.SERVICE]
        cves = [n for n in self.nodes.values() if n.type == NodeType.CVE]
        exploits = [n for n in self.nodes.values() if n.type == NodeType.EXPLOIT]
        creds = [n for n in self.nodes.values() if n.type == NodeType.CREDENTIAL]
        shells = [n for n in self.nodes.values() if n.type == NodeType.SHELL]
        wafs = [n for n in self.nodes.values() if n.type == NodeType.WAF]

        # Risk score: weighted sum
        critical_cves = [c for c in cves if c.properties.get("cvss", 0) >= 9.0]
        high_cves = [c for c in cves if 7.0 <= c.properties.get("cvss", 0) < 9.0]

        risk = (
            len(critical_cves) * 10 +
            len(high_cves) * 5 +
            len(exploits) * 8 +
            len(creds) * 15 +
            len(shells) * 25 -
            len(wafs) * 3  # WAFs reduce score slightly
        )

        chains = self.find_attack_chains()

        return {
            "risk_score": min(risk, 100),
            "hosts": len(hosts),
            "ports": len(ports),
            "services": len(services),
            "cves": len(cves),
            "critical_cves": len(critical_cves),
            "exploits": len(exploits),
            "credentials": len(creds),
            "shells": len(shells),
            "wafs": len(wafs),
            "attack_chains": len(chains),
            "shortest_chain": len(chains[0]) if chains else 0,
        }

    def get_nodes_by_type(self, node_type: NodeType) -> List[Node]:
        """Get all nodes of a specific type."""
        return [n for n in self.nodes.values() if n.type == node_type]

    def get_neighbors(self, node_id: str) -> List[Tuple[Edge, Node]]:
        """Get all neighbors of a node (outgoing edges)."""
        result = []
        for edge in self._adjacency.get(node_id, []):
            target = self.nodes.get(edge.target_id)
            if target:
                result.append((edge, target))
        return result

    def get_path_to(self, target_node_id: str) -> Optional[List[str]]:
        """Find shortest path from root target to a specific node (BFS)."""
        from collections import deque
        start = f"target:{self.target}"
        if start == target_node_id:
            return [start]

        queue = deque([(start, [start])])
        visited = {start}
        while queue:
            current, path = queue.popleft()
            for edge in self._adjacency.get(current, []):
                if edge.target_id not in visited:
                    new_path = path + [edge.target_id]
                    if edge.target_id == target_node_id:
                        return new_path
                    visited.add(edge.target_id)
                    queue.append((edge.target_id, new_path))
        return None

    # --- Context Generation for AI ---

    def to_ai_context(self, max_lines: int = 50) -> str:
        """Generate a compact context string for the AI agent."""
        lines = [f"## Knowledge Graph — {self.target}\n"]

        score = self.get_attack_surface_score()
        lines.append(f"**Risk Score:** {score['risk_score']}/100")
        lines.append(f"**Surface:** {score['hosts']} hosts, {score['ports']} ports, "
                     f"{score['services']} services")
        lines.append(f"**Vulns:** {score['cves']} CVEs ({score['critical_cves']} critical), "
                     f"{score['exploits']} exploits")
        lines.append(f"**Access:** {score['credentials']} creds, {score['shells']} shells")
        if score['wafs'] > 0:
            lines.append(f"**Defense:** {score['wafs']} WAFs detected")
        lines.append("")

        # Attack chains
        chains = self.find_attack_chains()
        if chains:
            lines.append(f"### Top Attack Chains ({len(chains)} found)")
            for i, chain in enumerate(chains[:5], 1):
                readable = " → ".join(
                    self.nodes[n_id].label if n_id in self.nodes else n_id
                    for n_id in chain
                )
                lines.append(f"  {i}. {readable}")
            lines.append("")

        # Critical CVEs
        critical = [n for n in self.nodes.values()
                    if n.type == NodeType.CVE and n.properties.get("cvss", 0) >= 7.0]
        if critical:
            critical.sort(key=lambda n: n.properties.get("cvss", 0), reverse=True)
            lines.append("### High/Critical CVEs")
            for cve in critical[:10]:
                lines.append(f"  - {cve.label} (CVSS {cve.properties.get('cvss', '?')}): "
                             f"{cve.properties.get('description', '')[:80]}")
            lines.append("")

        # Credentials
        creds = self.get_nodes_by_type(NodeType.CREDENTIAL)
        if creds:
            lines.append("### Credentials Found")
            for c in creds[:10]:
                lines.append(f"  - {c.label}")
            lines.append("")

        return '\n'.join(lines[:max_lines])

    def to_report_narrative(self) -> str:
        """Generate an attack narrative for the report."""
        chains = self.find_attack_chains()
        if not chains:
            return "No complete attack chains were discovered during this engagement."

        lines = ["## Attack Narrative\n"]
        lines.append(f"The penetration test against **{self.target}** revealed "
                     f"**{len(chains)}** viable attack chain(s).\n")

        for i, chain in enumerate(chains[:3], 1):
            lines.append(f"### Attack Chain {i}")
            steps = []
            for j, node_id in enumerate(chain):
                node = self.nodes.get(node_id)
                if node:
                    if node.type == NodeType.TARGET:
                        steps.append(f"Starting from the target **{node.label}**")
                    elif node.type == NodeType.HOST:
                        steps.append(f"resolved to host **{node.label}**")
                    elif node.type == NodeType.PORT:
                        steps.append(f"port **{node.label}** was found open")
                    elif node.type == NodeType.SERVICE:
                        steps.append(f"running **{node.label}**")
                    elif node.type == NodeType.CVE:
                        cvss = node.properties.get('cvss', '?')
                        steps.append(f"vulnerable to **{node.label}** (CVSS {cvss})")
                    elif node.type == NodeType.EXPLOIT:
                        steps.append(f"exploitable via **{node.label}**")
                    elif node.type == NodeType.CREDENTIAL:
                        steps.append(f"yielding credentials **{node.label}**")
                    elif node.type == NodeType.SHELL:
                        steps.append(f"granting **{node.label}**")

            lines.append(", ".join(steps) + ".\n")

        return '\n'.join(lines)

    # --- Populate from EngagementState ---

    def populate_from_state(self, state) -> 'KnowledgeGraph':
        """Populate the KG from an EngagementState object."""
        # Hosts
        for host in state.discovered_hosts:
            self.add_host(host)

        # Ports
        for port_info in state.discovered_ports:
            self.add_port(
                host=port_info.get("host", self.target),
                port=port_info.get("port", 0),
                protocol=port_info.get("protocol", "tcp"),
                service=port_info.get("service", ""),
                version=port_info.get("version", ""),
            )

        # Vulns
        for vuln in state.discovered_vulns:
            cve_id = vuln.get("name", vuln.get("template", ""))
            if cve_id.startswith("CVE-"):
                self.add_cve(
                    host=vuln.get("host", self.target),
                    port=vuln.get("port", 0),
                    cve_id=cve_id,
                    cvss=vuln.get("cvss", 0.0),
                    severity=vuln.get("severity", "unknown"),
                    description=vuln.get("description", ""),
                )

        # Credentials
        for cred in state.discovered_creds:
            self.add_credential(
                service=cred.get("service", "unknown"),
                host=cred.get("host", self.target),
                username=cred.get("username", ""),
                password=cred.get("password", ""),
            )

        return self

    # --- Serialization ---

    def save(self, filepath: str):
        """Save the knowledge graph to JSON."""
        data = {
            "target": self.target,
            "created_at": self.created_at,
            "nodes": {n_id: {
                "id": n.id, "type": n.type.value, "label": n.label,
                "properties": n.properties, "source": n.source,
                "confidence": n.confidence,
            } for n_id, n in self.nodes.items()},
            "edges": [{
                "source": e.source_id, "target": e.target_id,
                "type": e.type.value, "weight": e.weight,
                "properties": e.properties,
            } for e in self.edges],
            "stats": self.get_attack_surface_score(),
        }
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    @classmethod
    def load(cls, filepath: str) -> 'KnowledgeGraph':
        """Load a knowledge graph from JSON."""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)

        kg = cls(data["target"])
        kg.created_at = data.get("created_at", "")

        # Reconstruct nodes
        for n_id, n_data in data.get("nodes", {}).items():
            node = Node(
                id=n_data["id"],
                type=NodeType(n_data["type"]),
                label=n_data["label"],
                properties=n_data.get("properties", {}),
                source=n_data.get("source", ""),
                confidence=n_data.get("confidence", 1.0),
            )
            kg.nodes[n_id] = node

        # Reconstruct edges
        for e_data in data.get("edges", []):
            kg._add_edge(
                e_data["source"], e_data["target"],
                EdgeType(e_data["type"]),
                weight=e_data.get("weight", 1.0),
                **e_data.get("properties", {}),
            )

        return kg

    def __repr__(self):
        score = self.get_attack_surface_score()
        return (f"KnowledgeGraph({self.target}: {len(self.nodes)} nodes, "
                f"{len(self.edges)} edges, risk={score['risk_score']})")
