Update app.py

app.py

@@ -4,8 +4,13 @@ from transformers import pipeline
 import torch
 import time
 import re
+import logging
 
-# Fixed Professional Dashboard CSS - Complete Textbox Display
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Professional Dashboard CSS - Complete Textbox Display
 professional_css = """
 /* Professional SOC Dashboard - Fixed */
 .gradio-container {
@@ -321,88 +326,199 @@ model_status = "🔄 Loading..."
 
 @spaces.GPU
 def load_model():
-    """Load the best available model"""
+    """Load GPT-OSS-20B model with improved error handling"""
     global pipe, model_status
     
-    models_to_try = [
-        "openai/gpt-oss-20b",
-        "microsoft/DialoGPT-large",
-        "microsoft/DialoGPT-medium"
-    ]
-    
-    for model_name in models_to_try:
-        try:
-            pipe = pipeline(
-                "text-generation",
-                model=model_name,
-                torch_dtype="auto",
-                device_map="auto" if torch.cuda.is_available() else None,
-                trust_remote_code=True
-            )
-            pipe("Test", max_new_tokens=5, do_sample=False)
-            model_status = f"✅ {model_name.split('/')[-1]} Ready"
-            return model_status
-        except:
-            continue
-    
-    model_status = "⚠️ Fallback Mode"
-    return model_status
+    try:
+        logger.info("Starting model loading process...")
+        model_status = "🔄 Loading GPT-OSS-20B model..."
+        
+        # Load the specific model requested
+        logger.info("Loading gpt-oss-20b model...")
+        pipe = pipeline(
+            "text-generation",
+            model="openai/gpt-oss-20b",
+            torch_dtype=torch.float16,  # Use fp16 for better memory efficiency
+            device_map="auto",
+            trust_remote_code=True,
+            max_length=512,  # Limit context length
+            pad_token_id=50256  # Set pad token
+        )
+        
+        # Test the model with a simple prompt
+        logger.info("Testing model functionality...")
+        test_output = pipe(
+            "Test security analysis:", 
+            max_new_tokens=10, 
+            do_sample=True, 
+            temperature=0.7,
+            pad_token_id=50256
+        )
+        
+        model_status = "✅ GPT-OSS-20B Ready"
+        logger.info("Model loaded successfully!")
+        return model_status
+        
+    except Exception as e:
+        logger.error(f"Model loading failed: {str(e)}")
+        model_status = "⚠️ Model Loading Failed - Using Fallback"
+        pipe = None
+        return model_status
 
 @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()
     
-    # Quick pattern-based detection for demo
-    threats = []
-    if re.search(r'failed.*login|authentication.*failed', logs, re.IGNORECASE):
-        threats.append("🚨 Brute Force Attack")
-    if re.search(r'powershell.*-enc|cmd\.exe', logs, re.IGNORECASE):
-        threats.append("🚨 Malicious Script Execution")
-    if re.search(r'suspicious.*ip|unusual.*connection', logs, re.IGNORECASE):
-        threats.append("🚨 Suspicious Network Activity")
-    
-    if threats:
-        result = f"""🚨 THREATS DETECTED
+    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
         
-DETECTED THREATS:
-{chr(10).join(threats)}
+        # 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
 
-SEVERITY: {"Critical" if len(threats) > 2 else "High"}
-CONFIDENCE: {85 + len(threats) * 5}%
+AI ANALYSIS:
+{ai_analysis}
 
-IMMEDIATE ACTIONS:
-• Isolate affected systems
-• Preserve evidence
-• Escalate to L2 analyst
-• Implement containment"""
-        status = "🚨 THREATS DETECTED"
-    else:
-        result = """✅ NO THREATS DETECTED
+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" 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 = 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()
     
-    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:
@@ -414,8 +530,8 @@ IMMEDIATE ACTIONS:
 DECISION: Escalate to L2
 PRIORITY: High""",
 
-        "L2": f"""🔍 L2 INVESTIGATION
-        
+                "L2": f"""🔍 L2 INVESTIGATION
+                
 INCIDENT: {threat[:60]}...
 
 INVESTIGATION PLAN:
@@ -427,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:
@@ -439,19 +555,26 @@ 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
-SAMPLE_LOGS = """2025-08-12 14:30:15 [AUTH] Failed login: 'admin' from 192.168.1.100
-2025-08-12 14:30:18 [AUTH] Failed login: 'administrator' from 192.168.1.100  
-2025-08-12 14:30:45 [PROC] powershell.exe -WindowStyle Hidden -enc ZXhlYyBjYWxjLmV4ZQ==
-2025-08-12 14:31:12 [NET] Suspicious connection to 45.33.22.11:443"""
+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==
+2025-08-11 14:31:12 [NET] Suspicious connection to 45.33.22.11:443
+2025-08-11 14:31:30 [FILE] Unusual file access pattern detected
+2025-08-11 14:32:01 [NET] Multiple connections from same source IP"""
 
-SAMPLE_THREAT = "Multiple failed login attempts followed by encoded PowerShell execution and suspicious network traffic to external IP addresses."
+SAMPLE_THREAT = "Multiple failed login attempts detected from IP 192.168.1.100, followed by encoded PowerShell execution and suspicious outbound network connections to known malicious IP addresses. Lateral movement indicators present."
 
 # Main Dashboard Interface
 with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=professional_css) as demo:
@@ -460,14 +583,14 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
     gr.HTML("""
     <div class="dashboard-header">
         <div class="header-title">🛡️ SOC LLM Dashboard</div>
-        <div class="header-subtitle">Professional Security Operations Center • LLM-Powered Detection & Analysis</div>
+        <div class="header-subtitle">Professional Security Operations Center • GPT-OSS-20B Powered Detection & Analysis</div>
     </div>
     """)
     
     # System Status Bar
     with gr.Row():
         system_status = gr.Textbox(
-            value="🔄 Initializing AI Models...",
+            value="🔄 Initializing GPT-OSS-20B...",
             label="System Status",
             interactive=False,
             elem_classes=["status-indicator", "status-warning"],
@@ -480,7 +603,7 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
         
         # ================== TASK 1: DETECTION PANEL ==================
         with gr.Column(scale=1, elem_classes=["task-panel"]):
-            gr.HTML('<div class="task-header">📊 TASK 1: THREAT DETECTION</div>')
+            gr.HTML('<div class="task-header">📊 TASK 1: AI THREAT DETECTION</div>')
             
             # Detection Controls
             gr.HTML('<div class="control-label">Detection Sensitivity</div>')
@@ -492,13 +615,13 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
             )
             
             with gr.Row():
-                detect_btn = gr.Button("🔍 Detect", elem_classes=["primary-btn"], scale=2)
+                detect_btn = gr.Button("🔍 AI Detect", elem_classes=["primary-btn"], scale=2)
                 sample_logs_btn = gr.Button("📝 Sample", elem_classes=["secondary-btn"], scale=1)
             
             # Log Input
             gr.HTML('<div class="result-header">Security Logs Input</div>')
             log_input = gr.Textbox(
-                placeholder="Paste security logs here...",
+                placeholder="Paste security logs here for AI-powered analysis...",
                 lines=6,
                 elem_classes=["compact-input", "detection-input"],
                 interactive=True,
@@ -506,12 +629,12 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
             )
             
             # Detection Results
-            gr.HTML('<div class="result-header">Detection Results</div>')
+            gr.HTML('<div class="result-header">AI Detection Results</div>')
             detection_output = gr.Textbox(
                 lines=8,
                 elem_classes=["compact-output"],
                 interactive=False,
-                placeholder="Detection results will appear here...",
+                placeholder="GPT-OSS-20B detection results will appear here...",
                 show_label=False
             )
             
@@ -524,7 +647,7 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
         
         # ================== TASK 2: ASSISTANT PANEL ==================
         with gr.Column(scale=1, elem_classes=["task-panel"]):
-            gr.HTML('<div class="task-header">🤖 TASK 2: ANALYST ASSISTANT</div>')
+            gr.HTML('<div class="task-header">🤖 TASK 2: AI ANALYST ASSISTANT</div>')
             
             # Assistant Controls  
             gr.HTML('<div class="control-label">Analyst Level</div>')
@@ -536,13 +659,13 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
             )
             
             with gr.Row():
-                analyze_btn = gr.Button("🚀 Analyze", elem_classes=["primary-btn"], scale=2)
+                analyze_btn = gr.Button("🚀 AI Analyze", elem_classes=["primary-btn"], scale=2)
                 sample_threat_btn = gr.Button("📝 Sample", elem_classes=["secondary-btn"], scale=1)
             
             # Threat Input
             gr.HTML('<div class="result-header">Threat Description</div>')
             threat_input = gr.Textbox(
-                placeholder="Describe the security threat or incident...",
+                placeholder="Describe the security threat for AI analysis...",
                 lines=6,
                 elem_classes=["compact-input"],
                 interactive=True,
@@ -555,7 +678,7 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
                 lines=8,
                 elem_classes=["compact-output"],
                 interactive=False,
-                placeholder="Analysis results will appear here...",
+                placeholder="GPT-OSS-20B analysis results will appear here...",
                 show_label=False
             )
             
@@ -569,7 +692,7 @@ with gr.Blocks(title="SOC LLM Dashboard", theme=gr.themes.Soft(), css=profession
     # Quick Info Footer
     gr.HTML("""
     <div style="text-align: center; padding: 12px; color: rgba(255,255,255,0.8); font-size: 11px; margin-top: 10px;">
-        <strong>Research Project:</strong> LLM-based SOC Assistant • <strong>Student:</strong> Abdullah Alanazi • <strong>Supervisor:</strong> Prof. Ali Shoker • <strong>Institution:</strong> KAUST
+        <strong>Research Project:</strong> LLM-based SOC Assistant • <strong>Model:</strong> GPT-OSS-20B • <strong>Student:</strong> Abdullah Alanazi • <strong>Supervisor:</strong> Prof. Ali Shoker • <strong>Institution:</strong> KAUST
     </div>
     """)