DSATP and LlamaIndex with real CVE/IoT corpus

app.py

@@ -1,7 +1,12 @@
 import gradio as gr
 import yara
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 import os
 
+# Configure LlamaIndex to use local embeddings
+Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
+
 # DSATP log parsing (embedded for Spaces)
 def dsatp_parse_log(text: str) -> dict:
     """Parse log for IoT threats."""
@@ -40,9 +45,22 @@ def dsatp_yara_scan(file_path: str) -> dict:
     except Exception as e:
         return {"error": str(e), "severity": "Unknown", "mitigation": "Check file format"}
 
+# Initialize LlamaIndex with real corpus
+def init_llama_index():
+    """Load CVE/IoT corpus into LlamaIndex."""
+    try:
+        documents = SimpleDirectoryReader("corpus", filename_as_id=True).load_data()
+        return VectorStoreIndex.from_documents(documents)
+    except Exception as e:
+        print(f"Error loading corpus: {e}")
+        return None
+
+index = init_llama_index()
+query_engine = index.as_query_engine() if index else None
+
 # Chatbot function
 def chatbot_response(user_input, file, history):
-    """Process input or file with DSATP."""
+    """Process input or file with DSATP and LlamaIndex."""
     if history is None:
         history = []
     input_text = user_input
@@ -54,7 +72,16 @@ def chatbot_response(user_input, file, history):
     else:
         scan_result = dsatp_parse_log(input_text)
 
-    response = f"Security Analyst: {scan_result['classification']}. Severity: {scan_result['severity']}. Mitigation: {scan_result['mitigation']}. Confidence: {scan_result['confidence']:.1f}"
+    # Query LlamaIndex for context
+    context_str = "No context available."
+    if query_engine:
+        try:
+            context = query_engine.query(f"Mitigation for: {input_text}")
+            context_str = str(context)
+        except Exception as e:
+            context_str = f"Context error: {e}"
+
+    response = f"Security Analyst: {scan_result['classification']}. Severity: {scan_result['severity']}. Mitigation: {scan_result['mitigation']}. Confidence: {scan_result['confidence']:.1f}. Context: {context_str}"
     updated_history = history + [(user_input or "File uploaded", response)]
     return updated_history, scan_result
 

requirements.txt

@@ -3,4 +3,7 @@ textblob
 fastapi
 uvicorn
 yara-python
-requests
+requests
+llama-index-core
+pandas
+sentence-transformers