Update app.py

app.py

@@ -3,34 +3,57 @@ import spaces
 from transformers import pipeline
 import torch
 import time
+import re
+from datetime import datetime
 
-# Simple CSS for clean design
-simple_css = """
+# Enhanced CSS for both tasks
+enhanced_css = """
 .gradio-container {
-    max-width: 900px !important;
+    max-width: 1200px !important;
     margin: 0 auto !important;
     font-family: 'Arial', sans-serif;
 }
 
-.threat-input {
+.task-tab {
+    background: linear-gradient(135deg, #667eea, #764ba2) !important;
+    color: white !important;
     border-radius: 8px !important;
-    border: 2px solid #e0e0e0 !important;
+    padding: 12px 20px !important;
+    font-weight: 600 !important;
+    margin: 5px !important;
+}
+
+.detection-input {
+    border-radius: 8px !important;
+    border: 2px solid #e74c3c !important;
     padding: 15px !important;
-    font-size: 14px !important;
+    font-family: 'Courier New', monospace !important;
+    background: #2d3436 !important;
+    color: #ddd !important;
 }
 
-.threat-input:focus {
-    border-color: #667eea !important;
+.assistant-input {
+    border-radius: 8px !important;
+    border: 2px solid #667eea !important;
+    padding: 15px !important;
 }
 
-.analyze-btn {
-    background: #667eea !important;
-    border: none !important;
+.threat-detected {
+    background: linear-gradient(135deg, #e74c3c, #c0392b) !important;
+    color: white !important;
+    padding: 15px !important;
     border-radius: 8px !important;
-    padding: 12px 30px !important;
-    font-size: 16px !important;
+    margin: 10px 0 !important;
     font-weight: 600 !important;
+}
+
+.threat-clear {
+    background: linear-gradient(135deg, #27ae60, #229954) !important;
     color: white !important;
+    padding: 15px !important;
+    border-radius: 8px !important;
+    margin: 10px 0 !important;
+    font-weight: 600 !important;
 }
 
 .analysis-output {
@@ -41,13 +64,12 @@ simple_css = """
     line-height: 1.6 !important;
 }
 
-.status-box {
+.status-success {
     background: #d4edda !important;
     border: 1px solid #c3e6cb !important;
     color: #155724 !important;
     padding: 10px !important;
     border-radius: 6px !important;
-    margin: 10px 0 !important;
 }
 """
 
@@ -95,127 +117,341 @@ def load_model():
     model_status = "⚠️ Using fallback mode"
     return model_status
 
+# ===================== TASK 1: DETECTION =====================
+
+@spaces.GPU
+def detect_threats_from_logs(log_data, detection_sensitivity):
+    """Task 1: Detect threats from raw log data using LLM"""
+    
+    if not log_data.strip():
+        return "Please provide log data for analysis.", ""
+    
+    start_time = time.time()
+    
+    # Enhanced detection prompt
+    detection_prompt = f"""You are a cybersecurity threat detection system. Analyze these security logs and detect any potential threats, anomalies, or suspicious activities.
+
+DETECTION SENSITIVITY: {detection_sensitivity}
+
+RAW LOG DATA:
+{log_data}
+
+Analyze for:
+- Failed authentication attempts
+- Unusual network connections
+- Suspicious process executions
+- Privilege escalations
+- Malware indicators
+- Data exfiltration attempts
+- System anomalies
+
+Output format:
+THREAT DETECTED: [Yes/No]
+THREAT TYPE: [Type if detected]
+SEVERITY: [Critical/High/Medium/Low]
+CONFIDENCE: [Percentage]
+DETAILS: [Explanation]
+RECOMMENDATIONS: [Next steps]
+
+DETECTION ANALYSIS:"""
+
+    if pipe:
+        try:
+            result = pipe(
+                detection_prompt,
+                max_new_tokens=400,
+                do_sample=True,
+                temperature=0.2,  # Lower temperature for detection accuracy
+                top_p=0.9,
+                repetition_penalty=1.1
+            )
+            
+            detection_result = result[0]['generated_text'][len(detection_prompt):].strip()
+            
+            if len(detection_result) < 30:
+                detection_result = get_detection_fallback(log_data, detection_sensitivity)
+                
+        except Exception as e:
+            detection_result = f"AI Detection Error: {str(e)[:100]}\n\n{get_detection_fallback(log_data, detection_sensitivity)}"
+    else:
+        detection_result = get_detection_fallback(log_data, detection_sensitivity)
+    
+    # Determine if threat was detected for status display
+    threat_status = "🚨 THREAT DETECTED" if "THREAT DETECTED: Yes" in detection_result or "threat" in detection_result.lower() else "✅ NO THREATS DETECTED"
+    
+    processing_time = round(time.time() - start_time, 2)
+    status = f"{threat_status} | Analysis completed in {processing_time}s | {model_status}"
+    
+    return detection_result, status
+
+def get_detection_fallback(log_data, sensitivity):
+    """Fallback detection analysis using pattern matching"""
+    
+    # Simple pattern-based detection
+    threats_found = []
+    confidence = 60
+    
+    # Check for common threat indicators
+    if re.search(r'failed.*login|authentication.*failed|invalid.*password', log_data, re.IGNORECASE):
+        threats_found.append("Failed Authentication Attempts")
+        confidence += 20
+    
+    if re.search(r'powershell|cmd\.exe|suspicious.*process', log_data, re.IGNORECASE):
+        threats_found.append("Suspicious Process Execution")
+        confidence += 15
+    
+    if re.search(r'connection.*refused|unusual.*traffic|suspicious.*ip', log_data, re.IGNORECASE):
+        threats_found.append("Abnormal Network Activity")
+        confidence += 15
+    
+    if re.search(r'privilege.*escalation|admin.*rights|elevated.*access', log_data, re.IGNORECASE):
+        threats_found.append("Privilege Escalation Attempt")
+        confidence += 25
+    
+    if re.search(r'malware|virus|trojan|ransomware', log_data, re.IGNORECASE):
+        threats_found.append("Malware Indicators")
+        confidence += 30
+    
+    if threats_found:
+        severity = "Critical" if confidence > 85 else "High" if confidence > 70 else "Medium"
+        return f"""🚨 THREAT DETECTION ANALYSIS
+
+THREAT DETECTED: Yes
+THREAT TYPES: {', '.join(threats_found)}
+SEVERITY: {severity}
+CONFIDENCE: {min(confidence, 95)}%
+
+DETECTED INDICATORS:
+{chr(10).join(f"• {threat}" for threat in threats_found)}
+
+IMMEDIATE ACTIONS REQUIRED:
+• Isolate affected systems immediately
+• Preserve logs for forensic analysis
+• Escalate to L2 analyst for investigation
+• Implement containment measures
+• Monitor for lateral movement
+
+PATTERN ANALYSIS:
+Based on log pattern analysis, multiple threat indicators suggest ongoing malicious activity requiring immediate response."""
+
+    else:
+        return f"""✅ THREAT DETECTION ANALYSIS
+
+THREAT DETECTED: No
+SEVERITY: Low
+CONFIDENCE: {confidence}%
+
+ANALYSIS SUMMARY:
+No obvious threat indicators detected in the provided log data. However, this does not guarantee absence of sophisticated threats.
+
+RECOMMENDATIONS:
+• Continue monitoring for unusual patterns
+• Implement additional logging if needed
+• Consider advanced behavioral analysis
+• Regular security baseline reviews
+
+NOTE: Advanced persistent threats may use evasion techniques not detected by pattern analysis."""
+
+# ===================== TASK 2: ASSISTANT =====================
+
 @spaces.GPU
 def analyze_threat(threat_description, analyst_level):
-    """Simple threat analysis"""
+    """Task 2: Assist analysts with threat investigation"""
     
     if not threat_description.strip():
         return "Please enter a threat description first.", ""
     
     start_time = time.time()
     
-    # Create simple prompt
-    prompt = f"""As a {analyst_level} cybersecurity analyst, analyze this threat:
+    # Enhanced assistant prompt
+    assistant_prompt = f"""As a {analyst_level} cybersecurity analyst, provide detailed analysis for this security incident:
+
+INCIDENT: {threat_description}
 
-THREAT: {threat_description}
+Provide a comprehensive {analyst_level} level analysis including:
+- Threat assessment and classification
+- Potential impact and business risk
+- Investigation steps and evidence collection
+- Containment and mitigation strategies
+- Recommended actions and next steps
 
-Provide a {analyst_level} level security analysis including:
-- Threat assessment
-- Potential impact
-- Recommended actions
+Focus on {analyst_level} specific responsibilities and deliver actionable insights.
 
 ANALYSIS:"""
-    
+
     if pipe:
         try:
             result = pipe(
-                prompt,
-                max_new_tokens=300,
+                assistant_prompt,
+                max_new_tokens=400,
                 do_sample=True,
                 temperature=0.3,
                 top_p=0.9,
                 repetition_penalty=1.1
             )
             
-            analysis = result[0]['generated_text'][len(prompt):].strip()
+            analysis = result[0]['generated_text'][len(assistant_prompt):].strip()
             
             if len(analysis) < 30:
-                analysis = get_simple_fallback(threat_description, analyst_level)
+                analysis = get_assistant_fallback(threat_description, analyst_level)
                 
         except Exception as e:
-            analysis = f"AI Error: {str(e)[:100]}\n\n{get_simple_fallback(threat_description, analyst_level)}"
+            analysis = f"AI Analysis Error: {str(e)[:100]}\n\n{get_assistant_fallback(threat_description, analyst_level)}"
     else:
-        analysis = get_simple_fallback(threat_description, analyst_level)
+        analysis = get_assistant_fallback(threat_description, analyst_level)
     
     processing_time = round(time.time() - start_time, 2)
-    status = f"✅ Analysis completed in {processing_time}s | {model_status}"
+    status = f"✅ {analyst_level} analysis completed in {processing_time}s | {model_status}"
     
     return analysis, status
 
-def get_simple_fallback(threat_description, analyst_level):
-    """Simple fallback analysis"""
+def get_assistant_fallback(threat_description, analyst_level):
+    """Fallback assistant analysis"""
     
     if analyst_level == "L1":
         return f"""🚨 L1 TRIAGE ANALYSIS
 
-THREAT SUMMARY:
+INCIDENT SUMMARY:
 {threat_description}
 
-IMMEDIATE ACTIONS:
-• Assess severity and scope
-• Document all available evidence  
-• Isolate affected systems if needed
-• Escalate to L2 if severity is high
+IMMEDIATE TRIAGE ACTIONS:
+• Assess severity: Determine if this requires immediate escalation
+• Initial containment: Isolate affected systems if critical
+• Evidence preservation: Secure logs and system state
+• Documentation: Record all initial observations
+• Communication: Notify L2 analyst if severity warrants
+
+SEVERITY ASSESSMENT:
+• Impact scope: Determine number of affected systems
+• Data sensitivity: Assess if sensitive data is involved
+• Business criticality: Evaluate affected business functions
+• Time sensitivity: Determine urgency of response
 
-PRIORITY: Immediate containment and escalation decision required"""
+ESCALATION CRITERIA:
+• Critical/High severity incidents → Immediate L2 escalation
+• Multiple system involvement → L2 investigation required
+• Potential data breach → L2/L3 consultation needed
+• Advanced threat indicators → Expert analysis required
+
+NEXT STEPS:
+1. Complete initial assessment checklist
+2. Gather additional context if needed
+3. Make escalation decision based on severity
+4. Hand off to appropriate team with complete documentation"""
 
     elif analyst_level == "L2":
-        return f"""🔍 L2 INVESTIGATION ANALYSIS
+        return f"""🔍 L2 DETAILED INVESTIGATION
 
-THREAT DETAILS:
+INCIDENT DETAILS:
 {threat_description}
 
-INVESTIGATION STEPS:
-1. Collect and preserve evidence
-2. Analyze attack vectors and methods
-3. Determine scope of compromise
-4. Identify indicators of compromise (IOCs)
-5. Assess potential data exposure
+INVESTIGATION METHODOLOGY:
+1. Evidence Collection:
+   • System logs and event data
+   • Network traffic analysis
+   • File system artifacts
+   • Memory dumps if needed
+   • User activity records
 
-CONTAINMENT:
-• Implement network segmentation
-• Deploy additional monitoring
-• Review authentication logs
-• Check for lateral movement
+2. Timeline Analysis:
+   • Establish attack timeline
+   • Identify initial compromise vector
+   • Map lateral movement patterns
+   • Document persistence mechanisms
 
-NEXT STEPS:
-• Continue monitoring for related activity
-• Update security controls as needed
-• Consider L3 escalation for complex threats"""
+3. Technical Analysis:
+   • IOC identification and extraction
+   • Malware analysis if applicable
+   • Network communication analysis
+   • System configuration review
+
+4. Scope Assessment:
+   • Affected systems inventory
+   • Data exposure evaluation
+   • Privilege escalation review
+   • Lateral movement detection
+
+CONTAINMENT STRATEGY:
+• Network segmentation to prevent spread
+• Account restrictions for compromised users
+• System isolation for infected machines
+• Enhanced monitoring deployment
+
+RECOMMENDATIONS:
+• Deploy additional detection rules
+• Implement temporary compensating controls
+• Coordinate with infrastructure teams
+• Prepare for potential L3 escalation"""
 
     else:  # L3
-        return f"""🎯 L3 EXPERT ANALYSIS
+        return f"""🎯 L3 STRATEGIC ANALYSIS
 
 STRATEGIC THREAT ASSESSMENT:
 {threat_description}
 
-ADVANCED ANALYSIS:
-• Threat actor attribution assessment
-• Campaign analysis and TTPs
-• Business impact evaluation
-• Risk quantification
+EXECUTIVE SUMMARY:
+This incident requires senior-level analysis due to its complexity and potential business impact. Strategic coordination is needed for effective response.
 
-STRATEGIC RESPONSE:
-• Coordinate incident response team
-• Executive briefing preparation
-• Regulatory compliance review
-• Long-term security posture improvements
+THREAT LANDSCAPE ANALYSIS:
+• Adversary attribution and capability assessment
+• Campaign analysis and threat actor profiling
+• Attack sophistication and TTPs evaluation
+• Strategic intent and targeting analysis
 
-RECOMMENDATIONS:
-• Implement advanced threat hunting
-• Enhance detection capabilities
-• Review security architecture
-• Consider external forensics support"""
+BUSINESS IMPACT ASSESSMENT:
+• Operational disruption evaluation
+• Financial impact quantification
+• Regulatory compliance implications
+• Reputational risk assessment
+
+STRATEGIC RESPONSE PLAN:
+1. Crisis Management:
+   • Executive stakeholder notification
+   • Communication strategy development
+   • Resource allocation decisions
+   • External support engagement
+
+2. Advanced Investigation:
+   • Threat hunting operations
+   • Advanced forensics deployment
+   • Third-party security consultation
+   • Law enforcement coordination if needed
+
+3. Recovery Strategy:
+   • Business continuity planning
+   • System restoration priorities
+   • Security architecture review
+   • Lessons learned integration
+
+LONG-TERM RECOMMENDATIONS:
+• Security program enhancement
+• Advanced threat detection investment
+• Staff training and awareness programs
+• Strategic security partnerships"""
+
+# ===================== SAMPLE DATA =====================
+
+SAMPLE_LOGS = """2025-08-12 14:30:15 [AUTH] Failed login attempt for user 'administrator' from 192.168.1.100
+2025-08-12 14:30:18 [AUTH] Failed login attempt for user 'admin' from 192.168.1.100
+2025-08-12 14:30:22 [AUTH] Failed login attempt for user 'root' 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] Outbound connection to 45.33.22.11:443 from 192.168.1.100
+2025-08-12 14:31:45 [FILE] Suspicious file created: C:\\temp\\update.exe
+2025-08-12 14:32:10 [PROC] rundll32.exe comsvcs.dll MiniDump 1234 lsass.dmp full
+2025-08-12 14:32:33 [NET] Large data transfer detected: 1.2GB to external IP"""
 
-# Create simple interface
-with gr.Blocks(title="Simple SOC Analyzer", theme=gr.themes.Soft(), css=simple_css) as demo:
+SAMPLE_THREAT = "Suspicious PowerShell execution detected on user workstation with encoded commands, followed by unusual network traffic to external IP addresses and potential credential dumping activity."
+
+# ===================== GRADIO INTERFACE =====================
+
+with gr.Blocks(title="Comprehensive SOC LLM Tool", theme=gr.themes.Soft(), css=enhanced_css) as demo:
     
-    # Simple header
+    # Header
     gr.Markdown("""
-    # 🛡️ SOC Threat Analyzer
-    **Simple • Fast • Effective**
+    # 🛡️ Comprehensive SOC LLM Assistant
+    **Task 1: Threat Detection** • **Task 2: Analyst Assistant**
     
-    Enter any security threat and get instant AI analysis.
+    Covering both core LLM applications in Security Operations Centers
     """)
     
     # Model status
@@ -223,82 +459,186 @@ with gr.Blocks(title="Simple SOC Analyzer", theme=gr.themes.Soft(), css=simple_c
         value="🔄 Loading model...",
         label="System Status",
         interactive=False,
-        elem_classes=["status-box"]
+        elem_classes=["status-success"]
     )
     
-    # Main interface
-    with gr.Row():
-        with gr.Column(scale=1):
-            
-            # Threat input
-            threat_input = gr.Textbox(
-                label="🚨 Describe the Security Threat",
-                placeholder="Example: Suspicious PowerShell execution detected on user workstation with encoded commands...",
-                lines=5,
-                elem_classes=["threat-input"]
-            )
-            
-            # Analysis level
-            analyst_level = gr.Radio(
-                choices=["L1", "L2", "L3"],
-                value="L2",
-                label="Analysis Level",
-                info="L1: Quick Triage • L2: Detailed Investigation • L3: Strategic Analysis"
-            )
-            
-            # Analyze button
-            analyze_btn = gr.Button(
-                "🔍 Analyze Threat",
-                variant="primary",
-                size="lg",
-                elem_classes=["analyze-btn"]
-            )
-            
-            # Quick examples
+    # Tabbed interface for two tasks
+    with gr.Tabs():
+        
+        # ===================== TAB 1: DETECTION =====================
+        with gr.Tab("🔍 Task 1: Threat Detection", elem_classes=["task-tab"]):
             gr.Markdown("""
-            ### 📝 Quick Examples:
-            - Suspicious email with malicious attachment
-            - Unusual network traffic to external IP
-            - User account showing signs of compromise
-            - Ransomware indicators detected on server
-            - Failed login attempts from multiple locations
+            ## 🚨 Raw Log Analysis & Threat Detection
+            Upload or paste security logs to detect potential threats using LLM analysis.
             """)
-        
-        with gr.Column(scale=2):
             
-            # Analysis output
-            analysis_output = gr.Textbox(
-                label="🤖 Security Analysis",
-                lines=20,
-                interactive=False,
-                elem_classes=["analysis-output"],
-                placeholder="Analysis will appear here..."
-            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    
+                    # Detection controls
+                    detection_sensitivity = gr.Radio(
+                        choices=["High", "Medium", "Low"],
+                        value="Medium",
+                        label="🎯 Detection Sensitivity",
+                        info="High: More alerts, Low: Only obvious threats"
+                    )
+                    
+                    # Sample data button
+                    load_sample_btn = gr.Button(
+                        "📝 Load Sample Logs",
+                        variant="secondary"
+                    )
+                    
+                    detect_btn = gr.Button(
+                        "🔍 Detect Threats",
+                        variant="primary",
+                        size="lg"
+                    )
+                    
+                    gr.Markdown("""
+                    ### 📊 Detection Capabilities:
+                    - Failed authentication attempts
+                    - Suspicious process execution
+                    - Unusual network connections
+                    - Privilege escalation
+                    - Malware indicators
+                    - Data exfiltration patterns
+                    """)
+                
+                with gr.Column(scale=2):
+                    
+                    # Log input
+                    log_input = gr.Textbox(
+                        label="📋 Security Logs / System Events",
+                        placeholder="Paste your security logs here...\n\nExample:\n2025-08-12 14:30:15 [AUTH] Failed login attempt...\n2025-08-12 14:30:45 [PROC] powershell.exe -enc ...",
+                        lines=12,
+                        elem_classes=["detection-input"]
+                    )
+                    
+                    # Detection results
+                    detection_output = gr.Textbox(
+                        label="🚨 Threat Detection Results",
+                        lines=15,
+                        interactive=False,
+                        elem_classes=["analysis-output"],
+                        placeholder="Detection results will appear here..."
+                    )
+                    
+                    detection_status = gr.Textbox(
+                        label="Detection Status",
+                        interactive=False,
+                        lines=1
+                    )
+        
+        # ===================== TAB 2: ASSISTANT =====================
+        with gr.Tab("🤖 Task 2: Analyst Assistant", elem_classes=["task-tab"]):
+            gr.Markdown("""
+            ## 👥 Multi-Level Analyst Support
+            Get AI-powered analysis tailored to your expertise level (L1/L2/L3).
+            """)
             
-            # Processing status
-            process_status = gr.Textbox(
-                label="Processing Status",
-                interactive=False,
-                lines=1
-            )
-    
-    # Quick action buttons
-    with gr.Row():
-        gr.Button("💾 Save Analysis", variant="secondary", size="sm")
-        gr.Button("📧 Email Report", variant="secondary", size="sm") 
-        gr.Button("🔄 Clear All", variant="secondary", size="sm")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    
+                    # Assistant controls
+                    analyst_level = gr.Radio(
+                        choices=["L1", "L2", "L3"],
+                        value="L2",
+                        label="👤 Analyst Level",
+                        info="L1: Triage • L2: Investigation • L3: Strategic"
+                    )
+                    
+                    # Sample threat button
+                    load_threat_btn = gr.Button(
+                        "📝 Load Sample Threat",
+                        variant="secondary"
+                    )
+                    
+                    analyze_btn = gr.Button(
+                        "🚀 Analyze Threat",
+                        variant="primary",
+                        size="lg"
+                    )
+                    
+                    gr.Markdown("""
+                    ### 🎯 Analysis Levels:
+                    
+                    **L1 (Triage):**
+                    - Initial assessment
+                    - Containment actions
+                    - Escalation decisions
+                    
+                    **L2 (Investigation):**
+                    - Detailed analysis
+                    - Evidence collection
+                    - Technical investigation
+                    
+                    **L3 (Expert):**
+                    - Strategic assessment
+                    - Business impact
+                    - Executive briefing
+                    """)
+                
+                with gr.Column(scale=2):
+                    
+                    # Threat input
+                    threat_input = gr.Textbox(
+                        label="🚨 Threat Description",
+                        placeholder="Describe the security incident or threat...\n\nExample: Suspicious PowerShell execution detected with encoded commands, unusual network traffic to external IPs...",
+                        lines=8,
+                        elem_classes=["assistant-input"]
+                    )
+                    
+                    # Analysis results
+                    analysis_output = gr.Textbox(
+                        label="🤖 AI Analysis & Recommendations",
+                        lines=15,
+                        interactive=False,
+                        elem_classes=["analysis-output"],
+                        placeholder="Analysis will appear here..."
+                    )
+                    
+                    analysis_status = gr.Textbox(
+                        label="Analysis Status",
+                        interactive=False,
+                        lines=1
+                    )
     
-    # Simple footer
+    # Footer
     gr.Markdown("""
     ---
-    **💡 Tips:** Be specific about what you observed, include timestamps, IP addresses, user accounts, or file names when available.
+    ## 📊 **Comprehensive LLM-SOC Integration**
+    
+    **Task 1 (Detection):** Raw logs → Threat identification  
+    **Task 2 (Assistant):** Threat description → Investigation guidance
+    
+    *Demonstrating the full spectrum of LLM applications in Security Operations Centers*
     """)
     
-    # Event handlers
+    # ===================== EVENT HANDLERS =====================
+    
+    # Detection tab handlers
+    detect_btn.click(
+        fn=detect_threats_from_logs,
+        inputs=[log_input, detection_sensitivity],
+        outputs=[detection_output, detection_status]
+    )
+    
+    load_sample_btn.click(
+        fn=lambda: SAMPLE_LOGS,
+        outputs=[log_input]
+    )
+    
+    # Assistant tab handlers
     analyze_btn.click(
         fn=analyze_threat,
         inputs=[threat_input, analyst_level],
-        outputs=[analysis_output, process_status]
+        outputs=[analysis_output, analysis_status]
+    )
+    
+    load_threat_btn.click(
+        fn=lambda: SAMPLE_THREAT,
+        outputs=[threat_input]
     )
     
     # Initialize model on startup