Blame for code/graph/real/foundational.py

"""Load bounded, high-value public-world graph sources.

These loaders deliberately preserve source semantics.  ATT&CK and NVD edges are
security knowledge, not causal claims; GeoNames edges are geographic structure;
and Exoplanet Archive rows are astronomical observations.  Keeping the source
kind explicit prevents a mixed corpus from teaching the model that every edge
means cause and effect.
"""

from __future__ import annotations

import csv
import gzip
import json
import re
import zipfile
from collections.abc import Mapping
from pathlib import Path
from urllib.parse import unquote

from graph.graph import EdgeRecord
from graph.real.causal import _make_graph, _node_id
from graph.real.corpus import RealGraph, SourceMetadata, relation_phrase


def _short(value: object, limit: int = 4000) -> str:
    return " ".join(str(value or "").split())[:limit]


def _stix_identifier(value: str) -> str:
    return "stix:" + re.sub(r"[^A-Za-z0-9_.:-]+", "_", value)


def load_mitre_attack_stix(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_objects: int | None = None,
) -> RealGraph:
    """Load named ATT&CK STIX objects and their relationships."""
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    objects = payload.get("objects", [])
    if not isinstance(objects, list):
        raise TypeError("ATT&CK STIX bundle has no object list")

    selected: dict[str, Mapping[str, object]] = {}
    for item in objects:
        if not isinstance(item, dict):
            continue
        object_id = str(item.get("id", ""))
        object_type = str(item.get("type", ""))
        if not object_id or object_type in {"relationship", "marking-definition"}:
            continue
        if not item.get("name") and object_type not in {"identity"}:
            continue
        selected[object_id] = item
        if max_objects is not None and len(selected) >= max_objects:
            break

    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    for object_id, item in selected.items():
        node_id = _stix_identifier(object_id)
        object_type = str(item.get("type", "stix_object"))
        nodes[node_id] = object_type
        labels[node_id] = _short(item.get("name") or object_id, 240)
        external_ids = []
        references = item.get("external_references", ())
        if not isinstance(references, list):
            references = []
        for reference in references:
            if isinstance(reference, dict) and reference.get("external_id"):
                external_ids.append(str(reference["external_id"]))
        attributes[node_id] = {
            "stix_type": object_type,
            "description": _short(item.get("description")),
            "external_ids": ", ".join(external_ids[:16]),
            "evidence_kind": "mitre_attack_stix_object",
        }

    edges: list[EdgeRecord] = []
    relation_ids: set[str] = set()
    for item in objects:
        if not isinstance(item, dict) or item.get("type") != "relationship":
            continue
        left = str(item.get("source_ref", ""))
        right = str(item.get("target_ref", ""))
        if left not in selected or right not in selected or left == right:
            continue
        relation = (
            re.sub(
                r"[^a-z0-9_]+",
                "_",
                str(item.get("relationship_type", "related_to")).lower(),
            ).strip("_")
            or "related_to"
        )
        left_id = _stix_identifier(left)
        right_id = _stix_identifier(right)
        relation_ids.add(relation)
        edges.append(
            EdgeRecord(
                left_id,
                relation,
                right_id,
                nodes[left_id],
                nodes[right_id],
                relation_phrase(relation),
                {
                    "evidence_kind": "mitre_attack_stix_relationship",
                    "relationship_id": str(item.get("id", "")),
                    "description": _short(item.get("description"), 1000),
                },
            )
        )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids or {"related_to"},
    )


def _read_json(path: str | Path) -> Mapping[str, object]:
    file_path = Path(path)
    opener = gzip.open if file_path.name.endswith(".gz") else open
    with opener(file_path, "rt", encoding="utf-8") as handle:  # type: ignore[arg-type]
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise TypeError(f"expected JSON object in {file_path}")
    return payload


def _cpe_label(value: str) -> str:
    decoded = unquote(value)
    parts = decoded.split(":")
    if len(parts) >= 5:
        return ":".join(parts[:5])
    return decoded[:240]


def load_nvd_json(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_records: int | None = None,
) -> RealGraph:
    """Load NVD 2.0 CVE records, weaknesses, products, scores, and references."""
    payload = _read_json(path)
    raw_records = payload.get("vulnerabilities", payload.get("CVE_Items", ()))
    if not isinstance(raw_records, list):
        raise TypeError("NVD feed has no vulnerabilities list")
    records = raw_records if max_records is None else raw_records[:max_records]

    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relation_ids = {"has_weakness", "affects_product", "has_reference", "has_severity"}

    def add_node(node_id: str, node_type: str, label: str, **attrs: str) -> None:
        nodes[node_id] = node_type
        labels[node_id] = _short(label, 240)
        if attrs:
            attributes.setdefault(node_id, {}).update(attrs)

    def add_edge(
        left: str,
        relation: str,
        right: str,
        left_type: str,
        right_type: str,
        **attrs: str,
    ) -> None:
        edges.append(
            EdgeRecord(
                left,
                relation,
                right,
                left_type,
                right_type,
                relation_phrase(relation),
                {"evidence_kind": "nvd_cve_record", **attrs},
            )
        )

    for index, item in enumerate(records):
        if not isinstance(item, dict):
            continue
        cve = item.get("cve", item)
        if not isinstance(cve, dict):
            continue
        cve_id = str(
            cve.get("id")
            or cve.get("CVE_data_meta", {}).get("ID", f"CVE-unknown-{index}")
        )
        cve_node = f"nvd:cve:{cve_id}"
        descriptions = cve.get(
            "descriptions", cve.get("description", {}).get("description_data", ())
        )
        description = ""
        if isinstance(descriptions, list):
            for entry in descriptions:
                if isinstance(entry, dict) and (
                    not description or entry.get("lang") == "en"
                ):
                    description = _short(entry.get("value"))
                    if entry.get("lang") == "en":
                        break
        add_node(
            cve_node,
            "cve",
            cve_id,
            description=description,
            published=_short(cve.get("published")),
            last_modified=_short(cve.get("lastModified")),
        )

        weaknesses = cve.get("weaknesses", ())
        for weakness in weaknesses if isinstance(weaknesses, list) else ():
            if not isinstance(weakness, dict):
                continue
            descriptions = weakness.get("description", ())
            for entry in descriptions if isinstance(descriptions, list) else ():
                if not isinstance(entry, dict):
                    continue
                value = str(entry.get("value", "")).strip()
                if not value:
                    continue
                weakness_node = _node_id("nvd_cwe", value)
                add_node(weakness_node, "weakness", value)
                add_edge(cve_node, "has_weakness", weakness_node, "cve", "weakness")

        configurations = cve.get("configurations", ())
        stack: list[object] = (
            list(configurations) if isinstance(configurations, list) else []
        )
        while stack:
            current = stack.pop()
            if not isinstance(current, dict):
                continue
            stack.extend(
                current.get("children", ())
                if isinstance(current.get("children"), list)
                else ()
            )
            stack.extend(
                current.get("nodes", ())
                if isinstance(current.get("nodes"), list)
                else ()
            )
            matches = current.get("cpeMatch", ())
            for match in matches if isinstance(matches, list) else ():
                if not isinstance(match, dict):
                    continue
                cpe = str(match.get("criteria") or match.get("cpe23Uri") or "").strip()
                if not cpe:
                    continue
                product_node = _node_id("nvd_product", cpe)
                add_node(product_node, "affected_product", _cpe_label(cpe), cpe=cpe)
                add_edge(
                    cve_node, "affects_product", product_node, "cve", "affected_product"
                )

        references = cve.get("references", ())
        if isinstance(references, dict):
            references = references.get("reference_data", ())
        for reference in references if isinstance(references, list) else ():
            if not isinstance(reference, dict):
                continue
            url = str(reference.get("url", "")).strip()
            if not url:
                continue
            reference_node = _node_id("nvd_reference", url)
            add_node(
                reference_node,
                "reference",
                url,
                source=str(reference.get("source", "")),
            )
            add_edge(cve_node, "has_reference", reference_node, "cve", "reference")

        metrics = cve.get("metrics", {})
        if not isinstance(metrics, dict):
            metrics = {}
        score = severity = ""
        for metric_name in (
            "cvssMetricV40",
            "cvssMetricV31",
            "cvssMetricV30",
            "cvssMetricV2",
        ):
            entries = metrics.get(metric_name, ())
            if not isinstance(entries, list):
                continue
            for entry in entries:
                if not isinstance(entry, dict):
                    continue
                cvss = entry.get("cvssData", entry)
                if isinstance(cvss, dict):
                    score = str(cvss.get("baseScore", ""))
                    severity = str(cvss.get("baseSeverity", ""))
                if score or severity:
                    break
            if score or severity:
                break
        if score or severity:
            severity_node = _node_id("nvd_severity", f"{severity}-{score}")
            add_node(
                severity_node,
                "severity",
                f"{severity} {score}".strip(),
                score=score,
                severity=severity,
            )
            add_edge(cve_node, "has_severity", severity_node, "cve", "severity")

    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids,
    )


def load_nvd_cpe_dictionary(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_records: int | None = None,
) -> RealGraph:
    """Load the NVD CPE dictionary as a product/vendor/version ontology.

    CPE names are software/product identifiers, not vulnerability assertions.
    The adapter therefore emits product taxonomy edges and keeps the NVD CPE
    identifier in node attributes for auditability.
    """
    payload = _read_json(path)
    raw_products = payload.get("products", ())
    if not isinstance(raw_products, list):
        raise TypeError("NVD CPE feed has no products list")
    records = raw_products if max_records is None else raw_products[:max_records]
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relation_ids = {"has_vendor", "has_product", "has_version"}

    def add_node(node_id: str, node_type: str, label: str, **attrs: str) -> None:
        nodes[node_id] = node_type
        labels[node_id] = _short(label, 240)
        attributes.setdefault(node_id, {}).update(attrs)

    for item in records:
        if not isinstance(item, dict):
            continue
        cpe = item.get("cpe", item)
        if not isinstance(cpe, dict):
            continue
        uri = str(cpe.get("cpe23Uri") or cpe.get("cpeName") or "").strip()
        if not uri:
            continue
        parts = uri.split(":")
        if len(parts) < 6:
            continue
        vendor = unquote(parts[3]).replace("_", " ") or "unknown"
        product = unquote(parts[4]).replace("_", " ") or "unknown"
        version = unquote(parts[5]).replace("_", " ") or "unspecified"
        cpe_id = _node_id("nvd_cpe", uri)
        vendor_id = _node_id("nvd_vendor", vendor)
        product_id = _node_id("nvd_product_family", f"{vendor} {product}")
        version_id = _node_id("nvd_version", f"{vendor} {product} {version}")
        title = ""
        titles = cpe.get("titles", ())
        if isinstance(titles, list):
            for value in titles:
                if isinstance(value, dict) and value.get("title"):
                    title = str(value["title"])
                    if value.get("lang") in {"en", "en-us"}:
                        break
        add_node(cpe_id, "cpe_name", title or _cpe_label(uri), cpe_uri=uri)
        add_node(vendor_id, "vendor", vendor)
        add_node(product_id, "product_family", f"{vendor} {product}")
        add_node(version_id, "product_version", f"{product} {version}")
        edge_attrs = {"evidence_kind": "nvd_cpe_dictionary"}
        for left, relation, right, left_type, right_type in (
            (cpe_id, "has_vendor", vendor_id, "cpe_name", "vendor"),
            (cpe_id, "has_product", product_id, "cpe_name", "product_family"),
            (cpe_id, "has_version", version_id, "cpe_name", "product_version"),
        ):
            edges.append(
                EdgeRecord(
                    left,
                    relation,
                    right,
                    left_type,
                    right_type,
                    relation_phrase(relation),
                    edge_attrs,
                )
            )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids,
    )


def load_cisa_kev(path: str | Path, source: SourceMetadata) -> RealGraph:
    """Load CISA's Known Exploited Vulnerabilities catalog."""
    payload = _read_json(path)
    records = payload.get("vulnerabilities", ())
    if not isinstance(records, list):
        raise TypeError("CISA KEV feed has no vulnerabilities list")
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relation_ids = {"affects_product", "involves_vendor", "has_remediation", "due_by"}
    for item in records:
        if not isinstance(item, dict):
            continue
        cve = str(item.get("cveID", "")).strip()
        if not cve:
            continue
        cve_node = f"kev:cve:{cve}"
        nodes[cve_node] = "exploited_vulnerability"
        labels[cve_node] = cve
        attributes[cve_node] = {
            "vulnerability_name": _short(item.get("vulnerabilityName")),
            "description": _short(item.get("shortDescription")),
            "date_added": _short(item.get("dateAdded")),
        }
        for field, relation, node_type in (
            ("vendorProject", "involves_vendor", "vendor"),
            ("product", "affects_product", "product"),
            ("requiredAction", "has_remediation", "remediation"),
            ("dueDate", "due_by", "date"),
        ):
            value = str(item.get(field, "")).strip()
            if not value:
                continue
            target = _node_id("kev_" + node_type, value)
            nodes[target] = node_type
            labels[target] = _short(value, 240)
            edges.append(
                EdgeRecord(
                    cve_node,
                    relation,
                    target,
                    "exploited_vulnerability",
                    node_type,
                    relation_phrase(relation),
                    {"evidence_kind": "cisa_known_exploited_vulnerability"},
                )
            )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids,
    )


def _country_info(path: str | Path | None) -> dict[str, tuple[str, str]]:
    if path is None:
        return {}
    result: dict[str, tuple[str, str]] = {}
    with Path(path).open(encoding="utf-8") as handle:
        for line in handle:
            if not line.strip() or line.startswith("#"):
                continue
            fields = line.rstrip("\n").split("\t")
            if len(fields) > 8:
                result[fields[0]] = (fields[4], fields[8])
    return result


def load_geonames_cities(
    archive: str | Path,
    source: SourceMetadata,
    *,
    country_info: str | Path | None = None,
    max_rows: int | None = None,
) -> RealGraph:
    """Load a bounded GeoNames city/country/region/time-zone graph."""
    countries = _country_info(country_info)
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relation_ids = {"in_country", "in_region", "in_timezone", "within_continent"}
    with zipfile.ZipFile(archive) as handle:
        member = next(name for name in handle.namelist() if name.endswith(".txt"))
        with handle.open(member) as binary:
            reader = csv.reader(
                (line.decode("utf-8", "replace") for line in binary), delimiter="\t"
            )
            for index, fields in enumerate(reader):
                if max_rows is not None and index >= max_rows:
                    break
                if len(fields) < 18:
                    continue
                geoname_id, name, ascii_name = fields[:3]
                country_code, admin1, timezone = fields[8], fields[10], fields[17]
                if not geoname_id or not country_code:
                    continue
                city = f"geonames:city:{geoname_id}"
                country = f"geonames:country:{country_code}"
                region = _node_id("geonames_region", f"{country_code}:{admin1}")
                tz = _node_id("geonames_timezone", timezone) if timezone else ""
                nodes[city] = "city"
                labels[city] = ascii_name or name
                attributes[city] = {
                    "name": _short(name, 240),
                    "latitude": fields[4],
                    "longitude": fields[5],
                    "population": fields[14],
                    "feature_code": fields[7],
                }
                nodes[country] = "country"
                labels[country] = (
                    countries.get(country_code, (country_code, ""))[0] or country_code
                )
                edges.append(
                    EdgeRecord(
                        city,
                        "in_country",
                        country,
                        "city",
                        "country",
                        relation_phrase("in_country"),
                        {"evidence_kind": "geonames_city"},
                    )
                )
                if admin1:
                    nodes[region] = "admin_region"
                    labels[region] = f"{country_code} {admin1}"
                    edges.append(
                        EdgeRecord(
                            city,
                            "in_region",
                            region,
                            "city",
                            "admin_region",
                            relation_phrase("in_region"),
                            {"evidence_kind": "geonames_city"},
                        )
                    )
                if tz:
                    nodes[tz] = "timezone"
                    labels[tz] = timezone
                    edges.append(
                        EdgeRecord(
                            city,
                            "in_timezone",
                            tz,
                            "city",
                            "timezone",
                            relation_phrase("in_timezone"),
                            {"evidence_kind": "geonames_city"},
                        )
                    )
                continent = countries.get(country_code, ("", ""))[1]
                if continent:
                    continent_node = _node_id("geonames_continent", continent)
                    nodes[continent_node] = "continent"
                    labels[continent_node] = continent
                    edges.append(
                        EdgeRecord(
                            country,
                            "within_continent",
                            continent_node,
                            "country",
                            "continent",
                            relation_phrase("within_continent"),
                            {"evidence_kind": "geonames_country_info"},
                        )
                    )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids,
    )


def load_exoplanet_csv(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_rows: int | None = None,
) -> RealGraph:
    """Load confirmed-planet rows from NASA's Exoplanet Archive."""
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relation_ids = {"orbits", "discovered_by", "discovered_in", "has_measurement"}
    with Path(path).open(encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for index, row in enumerate(reader):
            if max_rows is not None and index >= max_rows:
                break
            planet_name = str(row.get("pl_name", "")).strip()
            host_name = str(row.get("hostname", "")).strip()
            if not planet_name or not host_name:
                continue
            planet = _node_id("exoplanet", planet_name)
            host = _node_id("exostar", host_name)
            nodes[planet] = "exoplanet"
            labels[planet] = planet_name
            nodes[host] = "host_star"
            labels[host] = host_name
            attributes[planet] = {
                "discovery_method": _short(row.get("discoverymethod")),
                "discovery_year": _short(row.get("disc_year")),
                "orbital_period_days": _short(row.get("pl_orbper")),
                "radius_earth": _short(row.get("pl_rade")),
                "mass_earth": _short(row.get("pl_bmasse")),
            }
            edges.append(
                EdgeRecord(
                    planet,
                    "orbits",
                    host,
                    "exoplanet",
                    "host_star",
                    relation_phrase("orbits"),
                    {"evidence_kind": "nasa_exoplanet_archive"},
                )
            )
            method = str(row.get("discoverymethod", "")).strip()
            if method:
                method_node = _node_id("exoplanet_method", method)
                nodes[method_node] = "discovery_method"
                labels[method_node] = method
                edges.append(
                    EdgeRecord(
                        planet,
                        "discovered_by",
                        method_node,
                        "exoplanet",
                        "discovery_method",
                        relation_phrase("discovered_by"),
                        {"evidence_kind": "nasa_exoplanet_archive"},
                    )
                )
            year = str(row.get("disc_year", "")).strip()
            if year:
                year_node = f"exoplanet:year:{year}"
                nodes[year_node] = "year"
                labels[year_node] = year
                edges.append(
                    EdgeRecord(
                        planet,
                        "discovered_in",
                        year_node,
                        "exoplanet",
                        "year",
                        relation_phrase("discovered_in"),
                        {"evidence_kind": "nasa_exoplanet_archive"},
                    )
                )
            for field, value in (
                ("pl_rade", row.get("pl_rade")),
                ("pl_bmasse", row.get("pl_bmasse")),
                ("st_teff", row.get("st_teff")),
            ):
                text = str(value or "").strip()
                if not text:
                    continue
                measure_node = _node_id("exoplanet_measure", f"{field}:{text}")
                nodes[measure_node] = "measurement"
                labels[measure_node] = f"{field} {text}"
                edges.append(
                    EdgeRecord(
                        planet,
                        "has_measurement",
                        measure_node,
                        "exoplanet",
                        "measurement",
                        relation_phrase("has_measurement"),
                        {"evidence_kind": "nasa_exoplanet_archive", "field": field},
                    )
                )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relation_ids,
    )