Update app.py

app.py

@@ -367,65 +367,158 @@ def load_model():
 
 @spaces.GPU
 def detect_threats(logs, sensitivity):
-    """Task 1: Threat Detection"""
+    """Task 1: AI-powered Threat Detection"""
+    global pipe
+    
     if not logs.strip():
         return "Please provide log data.", "⚠️ No input"
     
     start_time = time.time()
     
-    # Enhanced pattern-based detection that matches the screenshot
-    threats = []
-    
-    # Check for failed login attempts
-    if re.search(r'failed.*login|authentication.*failed', logs, re.IGNORECASE):
-        threats.append("🚨 Brute Force Attack")
-    
-    # Check for PowerShell execution
-    if re.search(r'powershell.*-enc|cmd\.exe', logs, re.IGNORECASE):
-        threats.append("🚨 Malicious Script Execution")
-    
-    # Check for suspicious network activity
-    if re.search(r'suspicious.*ip|unusual.*connection', logs, re.IGNORECASE):
-        threats.append("🚨 Suspicious Network Activity")
-    
-    # Generate result matching the original screenshot format
-    if threats:
-        result = f"""ASSESSMENT:
-• Risk Score: 70/100
-• Severity: CRITICAL
-• Confidence: 85%
-• Model: Pattern-based
+    try:
+        if pipe is not None:
+            # Use GPT-OSS-20B for AI-powered detection
+            prompt = f"""Analyze these security logs for threats:
+
+{logs}
+
+Detection sensitivity: {sensitivity}
+
+Analysis:"""
+
+            response = pipe(
+                prompt,
+                max_new_tokens=200,
+                do_sample=True,
+                temperature=0.3,
+                pad_token_id=50256,
+                truncation=True
+            )
+            
+            ai_analysis = response[0]['generated_text'].split("Analysis:")[-1].strip()
+            
+        else:
+            # Fallback to pattern-based detection
+            ai_analysis = "AI model unavailable. Using pattern-based detection."
+        
+        # Enhanced pattern-based detection as backup/supplement
+        threats = []
+        risk_score = 0
+        
+        # Authentication threats
+        failed_logins = len(re.findall(r'failed.*login|authentication.*failed', logs, re.IGNORECASE))
+        if failed_logins > 3:
+            threats.append(f"🚨 Brute Force Attack ({failed_logins} failed attempts)")
+            risk_score += 30
+        elif failed_logins > 0:
+            threats.append(f"⚠️ Failed Authentication ({failed_logins} attempts)")
+            risk_score += 15
+        
+        # Malicious execution
+        if re.search(r'powershell.*-enc|cmd\.exe|eval\(|exec\(', logs, re.IGNORECASE):
+            threats.append("🚨 Malicious Script Execution")
+            risk_score += 35
+        
+        # Network anomalies
+        if re.search(r'suspicious.*ip|unusual.*connection', logs, re.IGNORECASE):
+            threats.append("🚨 Suspicious Network Activity")
+            risk_score += 25
+        
+        # File anomalies
+        if re.search(r'unusual.*file|suspicious.*access', logs, re.IGNORECASE):
+            threats.append("⚠️ File System Anomaly")
+            risk_score += 20
+        
+        # Generate final result
+        if threats or pipe is not None:
+            severity = "CRITICAL" if risk_score > 50 else "HIGH" if risk_score > 30 else "MEDIUM"
+            confidence = min(95, 70 + len(threats) * 5)
+            
+            result = f"""🚨 THREAT ANALYSIS RESULTS
+
+AI ANALYSIS:
+{ai_analysis}
+
+DETECTED PATTERNS:
+{chr(10).join(f"• {threat}" for threat in threats) if threats else "• No obvious threat patterns detected"}
+
+ASSESSMENT:
+• Risk Score: {risk_score}/100
+• Severity: {severity if threats else "LOW"}
+• Confidence: {confidence}%
+• Model: {"GPT-OSS-20B" if pipe else "Pattern-based"}
 
 RECOMMENDATIONS:
-• Immediate containment required
-• Escalate to L2 analyst
+• {"Immediate containment required" if risk_score > 40 else "Continue monitoring"}
+• {"Escalate to L2 analyst" if risk_score > 30 else "Standard response"}
 • Preserve all evidence
 • Update threat intelligence"""
-        status = "🚨 THREATS DETECTED"
-    else:
-        result = """✅ NO THREATS DETECTED
+            
+            status = f"🚨 Analysis Complete - {len(threats)} threats found" if threats else "✅ Analysis Complete"
+        else:
+            result = """✅ NO THREATS DETECTED
+
+Clean log analysis with no suspicious patterns identified.
+Continue standard monitoring procedures."""
+            status = "✅ CLEAN"
         
-ANALYSIS: Clean logs
-CONFIDENCE: 75%
-STATUS: Normal operation
-RECOMMENDATION: Continue monitoring"""
-        status = "✅ CLEAN"
-    
-    time_taken = round(time.time() - start_time, 1)
-    return result, f"{status} ({time_taken}s)"
+        time_taken = round(time.time() - start_time, 1)
+        return result, f"{status} ({time_taken}s)"
+        
+    except Exception as e:
+        logger.error(f"Detection error: {str(e)}")
+        return f"❌ Analysis failed: {str(e)}", "❌ ERROR"
 
 @spaces.GPU
 def analyze_threat(threat, level):
-    """Task 2: Analyst Assistant"""
+    """Task 2: AI-powered Analyst Assistant"""
+    global pipe
+    
     if not threat.strip():
         return "Please describe the threat.", "⚠️ No input"
     
     start_time = time.time()
     
-    # Analysis templates that match the original screenshot format
-    templates = {
-        "L1": f"""🚨 L1 TRIAGE
-        
+    try:
+        if pipe is not None:
+            # Use GPT-OSS-20B for AI analysis
+            prompt = f"""As a Level {level} SOC analyst, analyze this security threat:
+
+{threat}
+
+Provide detailed analysis including:
+1. Threat assessment
+2. Recommended actions
+3. Priority level
+4. Next steps
+
+Analysis:"""
+
+            response = pipe(
+                prompt,
+                max_new_tokens=300,
+                do_sample=True,
+                temperature=0.4,
+                pad_token_id=50256,
+                truncation=True
+            )
+            
+            ai_analysis = response[0]['generated_text'].split("Analysis:")[-1].strip()
+            
+            result = f"""🤖 AI-POWERED {level} ANALYSIS
+
+THREAT ASSESSMENT:
+{ai_analysis}
+
+MODEL: GPT-OSS-20B
+ANALYST LEVEL: {level}
+STATUS: AI Analysis Complete"""
+            
+        else:
+            # Fallback analysis templates
+            templates = {
+                "L1": f"""🚨 L1 TRIAGE ANALYSIS
+                
 THREAT: {threat[:60]}...
 
 IMMEDIATE ACTIONS:
@@ -437,8 +530,8 @@ IMMEDIATE ACTIONS:
 DECISION: Escalate to L2
 PRIORITY: High""",
 
-        "L2": f"""🔍 L2 INVESTIGATION
-        
+                "L2": f"""🔍 L2 INVESTIGATION
+                
 INCIDENT: {threat[:60]}...
 
 INVESTIGATION PLAN:
@@ -450,8 +543,8 @@ INVESTIGATION PLAN:
 
 NEXT STEPS: Deploy monitoring""",
 
-        "L3": f"""🎯 L3 STRATEGIC ANALYSIS
-        
+                "L3": f"""🎯 L3 STRATEGIC ANALYSIS
+                
 THREAT ASSESSMENT: {threat[:60]}...
 
 STRATEGIC RESPONSE:
@@ -462,13 +555,18 @@ STRATEGIC RESPONSE:
 • Security improvements
 
 RECOMMENDATION: Full IR activation"""
-    }
-    
-    result = templates.get(level, templates["L2"])
-    time_taken = round(time.time() - start_time, 1)
-    return result, f"✅ {level} Complete ({time_taken}s)"
+            }
+            
+            result = templates.get(level, templates["L2"])
+        
+        time_taken = round(time.time() - start_time, 1)
+        return result, f"✅ {level} Complete ({time_taken}s)"
+        
+    except Exception as e:
+        logger.error(f"Analysis error: {str(e)}")
+        return f"❌ Analysis failed: {str(e)}", "❌ ERROR"
 
-# Sample data - matches the scenario in the screenshot
+# Sample data
 SAMPLE_LOGS = """2025-08-11 14:30:15 [AUTH] Failed login: 'admin' from 192.168.1.100
 2025-08-11 14:30:18 [AUTH] Failed login: 'administrator' from 192.168.1.100  
 2025-08-11 14:30:45 [PROC] powershell.exe -WindowStyle Hidden -enc ZXhlYyBjYWxjLmV4ZQ==