Spaces:
Configuration error
Configuration error
| """ | |
| 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.") | |