""" ingest.py Downloads NVD, CISA KEV, and Exploit-DB data, embeds it, and upserts to Pinecone. Run once locally before deploying. Usage: python scripts/ingest.py """ import os import sys import csv import json import time import requests from io import StringIO from tqdm import tqdm sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from pinecone import Pinecone, ServerlessSpec from sentence_transformers import SentenceTransformer from src.config import ( PINECONE_API_KEY, PINECONE_INDEX, PINECONE_ENV, EMBED_MODEL, EMBED_DIMENSION, NAMESPACE_NVD, NAMESPACE_KEV, NAMESPACE_EXPLOIT, NVD_BASE_URL, NVD_RESULTS_PER_PAGE, NVD_START_YEAR, NVD_END_YEAR, KEV_URL, EXPLOIT_DB_URL, CHUNK_SIZE, CHUNK_OVERLAP, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def get_index(): """Connect to Pinecone and return the index, creating it if needed.""" pc = Pinecone(api_key=PINECONE_API_KEY) # pinecone-client v3 uses .names() to get a list of index name strings existing_indexes = pc.list_indexes().names() if PINECONE_INDEX not in existing_indexes: print(f"Creating index: {PINECONE_INDEX}") pc.create_index( name=PINECONE_INDEX, dimension=EMBED_DIMENSION, metric="cosine", spec=ServerlessSpec(cloud="aws", region=PINECONE_ENV), ) while not pc.describe_index(PINECONE_INDEX).status["ready"]: time.sleep(2) else: print(f"Index '{PINECONE_INDEX}' already exists. Connecting...") return pc.Index(PINECONE_INDEX) def embed(model, texts): return model.encode(texts, normalize_embeddings=True).tolist() def upsert(index, vectors, namespace, batch=100): for i in range(0, len(vectors), batch): index.upsert(vectors=vectors[i:i + batch], namespace=namespace) def chunk_text(text, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP): words = text.split() chunks, start = [], 0 while start < len(words): chunks.append(" ".join(words[start:start + size])) start += size - overlap return chunks # --------------------------------------------------------------------------- # NVD # --------------------------------------------------------------------------- def ingest_nvd(index, model): print("\n--- Ingesting NVD ---") total = 0 for year in range(NVD_START_YEAR, NVD_END_YEAR + 1): print(f" Year {year}") start = 0 while True: resp = requests.get(NVD_BASE_URL, params={ "pubStartDate": f"{year}-01-01T00:00:00.000", "pubEndDate": f"{year}-12-31T23:59:59.999", "startIndex": start, "resultsPerPage": NVD_RESULTS_PER_PAGE, }, timeout=30) resp.raise_for_status() data = resp.json() vulns = data.get("vulnerabilities", []) if not vulns: break vectors = [] for item in tqdm(vulns, desc=" Embedding", leave=False): cve = item.get("cve", {}) cve_id = cve.get("id", "") descriptions = cve.get("descriptions", []) desc = next((d["value"] for d in descriptions if d.get("lang") == "en"), "") if not cve_id or not desc: continue metrics = cve.get("metrics", {}) cvss_data = ( metrics.get("cvssMetricV31", [{}])[0].get("cvssData", {}) or metrics.get("cvssMetricV30", [{}])[0].get("cvssData", {}) or metrics.get("cvssMetricV2", [{}])[0].get("cvssData", {}) ) score = cvss_data.get("baseScore", 0.0) severity = cvss_data.get("baseSeverity", "UNKNOWN") vector_s = cvss_data.get("vectorString", "") weaknesses = cve.get("weaknesses", []) cwe = weaknesses[0].get("description", [{}])[0].get("value", "") if weaknesses else "" full_text = ( f"CVE ID: {cve_id}\n" f"CVSS Score: {score} ({severity})\n" f"CVSS Vector: {vector_s}\n" f"CWE: {cwe}\n" f"Description: {desc}" ) for i, chunk in enumerate(chunk_text(full_text)): emb = embed(model, [chunk]) vectors.append({ "id": f"{cve_id}-chunk-{i}", "values": emb[0], "metadata": { "cve_id": cve_id, "cvss_score": float(score), "severity": severity, "cwe": cwe, "source": "nvd", "text": chunk, }, }) if vectors: upsert(index, vectors, NAMESPACE_NVD) total += len(vectors) start += NVD_RESULTS_PER_PAGE if start >= data.get("totalResults", 0): break time.sleep(1) # NVD rate limit: 5 req/30s without API key print(f" NVD done. {total} chunks indexed.") # --------------------------------------------------------------------------- # CISA KEV # --------------------------------------------------------------------------- def ingest_kev(index, model): print("\n--- Ingesting CISA KEV ---") data = requests.get(KEV_URL, timeout=30).json() vectors = [] for item in tqdm(data.get("vulnerabilities", []), desc=" Embedding"): cve_id = item.get("cveID", "") if not cve_id: continue full_text = ( f"CVE ID: {cve_id}\n" f"Vendor: {item.get('vendorProject','')} | Product: {item.get('product','')}\n" f"Name: {item.get('vulnerabilityName','')}\n" f"Date Added to KEV: {item.get('dateAdded','')}\n" f"Ransomware Use: {item.get('knownRansomwareCampaignUse','Unknown')}\n" f"Description: {item.get('shortDescription','')}\n" f"Required Action: {item.get('requiredAction','')}" ) emb = embed(model, [full_text]) vectors.append({ "id": f"kev-{cve_id}", "values": emb[0], "metadata": { "cve_id": cve_id, "vendor": item.get("vendorProject", ""), "date_added": item.get("dateAdded", ""), "ransomware": item.get("knownRansomwareCampaignUse", ""), "source": "cisa-kev", "text": full_text, }, }) upsert(index, vectors, NAMESPACE_KEV) print(f" KEV done. {len(vectors)} entries indexed.") # --------------------------------------------------------------------------- # Exploit-DB # --------------------------------------------------------------------------- def ingest_exploit_db(index, model): print("\n--- Ingesting Exploit-DB ---") resp = requests.get(EXPLOIT_DB_URL, timeout=30) vectors = [] for row in tqdm(csv.DictReader(StringIO(resp.text)), desc=" Embedding"): edb_id = row.get("id", "").strip() desc = row.get("description", "").strip() if not edb_id or not desc: continue full_text = ( f"Exploit-DB ID: {edb_id}\n" f"CVE Reference: {row.get('codes','').strip()}\n" f"Type: {row.get('type','').strip()} | Platform: {row.get('platform','').strip()}\n" f"Date: {row.get('date_published', row.get('date','')).strip()}\n" f"Description: {desc}" ) emb = embed(model, [full_text]) vectors.append({ "id": f"edb-{edb_id}", "values": emb[0], "metadata": { "edb_id": edb_id, "cve_ids": row.get("codes", "").strip(), "type": row.get("type", "").strip(), "source": "exploit-db", "text": full_text, }, }) upsert(index, vectors, NAMESPACE_EXPLOIT) print(f" Exploit-DB done. {len(vectors)} entries indexed.") # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": print("Loading embedding model...") model = SentenceTransformer(EMBED_MODEL) print("Connecting to Pinecone...") index = get_index() ingest_nvd(index, model) ingest_kev(index, model) ingest_exploit_db(index, model) print("\nIngestion complete.")