DSATP, LlamaIndex, and Mistral-7B Modal API integration

app.py

@@ -1,50 +1,21 @@
 import gradio as gr
-import yara
+import requests
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 import os
 
+# Suppress Hugging Face symlink warning
+os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
+
+# DSATP API endpoints
+DSATP_API = "http://localhost:8000"
+
+# Modal Mistral-7B API endpoint (replace with your Modal URL)
+MODAL_API = "https://rithvickkumar27--mistral-7b-api-analyze.modal.run"
+
 # 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."""
-    log = text.lower()
-    threats = {
-        "compromised": {"classification": "Threat Detected", "severity": "Critical", "mitigation": "Isolate process, run port scan"},
-        "unauthorized": {"classification": "Threat Detected", "severity": "High", "mitigation": "Quarantine MAC address"},
-        "high cpu": {"classification": "Threat Detected", "severity": "Medium", "mitigation": "Check for crypto-miner or DoS"},
-        "inbound traffic": {"classification": "Threat Detected", "severity": "Medium", "mitigation": "Block closed ports"},
-        "firmware mismatch": {"classification": "Threat Detected", "severity": "High", "mitigation": "Validate OTA or rollback"}
-    }
-    for key, value in threats.items():
-        if key in log:
-            return {**value, "confidence": 0.9}
-    return {"classification": "No Threat", "severity": "Safe", "mitigation": "None", "confidence": 0.5}
-
-# DSATP YARA scanning (embedded for Spaces)
-def dsatp_yara_scan(file_path: str) -> dict:
-    """Scan file with YARA rules."""
-    try:
-        rules = yara.compile(filepath="dsatp/rules.yar")
-        matches = rules.match(file_path)
-        if matches:
-            return {
-                "classification": "Malware Detected",
-                "severity": "Critical",
-                "mitigation": "Quarantine file, run antivirus",
-                "confidence": 0.95
-            }
-        return {
-            "classification": "No Malware",
-            "severity": "Safe",
-            "mitigation": "None",
-            "confidence": 0.7
-        }
-    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."""
@@ -56,11 +27,42 @@ def init_llama_index():
         return None
 
 index = init_llama_index()
-query_engine = index.as_query_engine() if index else None
+query_engine = index.as_retriever() if index else None
+
+# Call Modal Mistral-7B API
+def call_mistral_llm(prompt):
+    """Call Mistral-7B via Modal API."""
+    try:
+        response = requests.post(MODAL_API, json={"prompt": prompt}, timeout=30)
+        if response.status_code == 200:
+            print(f"LLM response: {response.json()}")
+            return response.json().get("response", "LLM error")
+        else:
+            return f"Modal API error: {response.status_code}"
+    except Exception as e:
+        return f"Modal error: {e}"
+
+# Call DSATP log parsing
+def dsatp_parse_log(text: str) -> dict:
+    """Call DSATP log parsing API."""
+    try:
+        response = requests.post(f"{DSATP_API}/parse_log", json={"log": text})
+        return response.json() if response.status_code == 200 else {"error": "API error"}
+    except Exception as e:
+        return {"error": str(e), "severity": "Unknown", "mitigation": "Check API"}
+
+# Call DSATP YARA scanning
+def dsatp_yara_scan(file_path: str) -> dict:
+    """Call DSATP YARA scanning API."""
+    try:
+        response = requests.post(f"{DSATP_API}/scan_file", json={"file_path": file_path})
+        return response.json() if response.status_code == 200 else {"error": "API error"}
+    except Exception as e:
+        return {"error": str(e), "severity": "Unknown", "mitigation": "Check API"}
 
 # Chatbot function
 def chatbot_response(user_input, file, history):
-    """Process input or file with DSATP and LlamaIndex."""
+    """Process input or file with DSATP, LlamaIndex, and Mistral-7B."""
     if history is None:
         history = []
     input_text = user_input
@@ -76,12 +78,24 @@ def chatbot_response(user_input, file, history):
     context_str = "No context available."
     if query_engine:
         try:
-            context = query_engine.query(f"Mitigation for: {input_text}")
-            context_str = str(context)
+            results = query_engine.retrieve(f"Mitigation for: {input_text}")
+            context_str = " ".join([res.text for res in results[:2]])
         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}"
+    # Generate response with Mistral-7B via Modal
+    prompt = f"""
+    You are a Security Analyst. Based on:
+    Classification: {scan_result['classification']}
+    Severity: {scan_result['severity']}
+    Mitigation: {scan_result['mitigation']}
+    Confidence: {scan_result['confidence']}
+    Context: {context_str}
+    Provide a concise response to the user.
+    """
+    llm_response = call_mistral_llm(prompt)
+    response = f"Security Analyst: {llm_response}. Context: {context_str}"
+
     updated_history = history + [(user_input or "File uploaded", response)]
     return updated_history, scan_result
 

requirements.txt

@@ -5,5 +5,8 @@ uvicorn
 yara-python
 requests
 llama-index-core
+llama-index-embeddings-huggingface
 pandas
-sentence-transformers
+sentence-transformers
+transformers
+torch