app.py

import gradio as gr
import spaces
from transformers import pipeline
import torch
import time

# Simple CSS for clean design
simple_css = """
.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
    font-family: 'Arial', sans-serif;
}

.threat-input {
    border-radius: 8px !important;
    border: 2px solid #e0e0e0 !important;
    padding: 15px !important;
    font-size: 14px !important;
}

.threat-input:focus {
    border-color: #667eea !important;
}

.analyze-btn {
    background: #667eea !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 12px 30px !important;
    font-size: 16px !important;
    font-weight: 600 !important;
    color: white !important;
}

.analysis-output {
    background: #f8f9fa !important;
    border-radius: 8px !important;
    border: 1px solid #e0e0e0 !important;
    padding: 20px !important;
    line-height: 1.6 !important;
}

.status-box {
    background: #d4edda !important;
    border: 1px solid #c3e6cb !important;
    color: #155724 !important;
    padding: 10px !important;
    border-radius: 6px !important;
    margin: 10px 0 !important;
}
"""

# Global model variables
pipe = None
model_status = "🔄 Loading model..."

@spaces.GPU
def load_model():
    """Load the best available model"""
    global pipe, model_status
    
    models_to_try = [
        "openai/gpt-oss-20b",
        "microsoft/DialoGPT-large",
        "microsoft/DialoGPT-medium",
        "gpt2-large"
    ]
    
    for model_name in models_to_try:
        try:
            print(f"🔄 Loading {model_name}...")
            
            pipe = pipeline(
                "text-generation",
                model=model_name,
                torch_dtype="auto",
                device_map="auto" if torch.cuda.is_available() else None,
                trust_remote_code=True
            )
            
            # Test the model
            pipe("Test", max_new_tokens=5, do_sample=False)
            
            model_status = f"✅ {model_name} ready"
            print(model_status)
            return model_status
            
        except Exception as e:
            print(f"❌ {model_name} failed: {str(e)[:50]}")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            continue
    
    model_status = "⚠️ Using fallback mode"
    return model_status

@spaces.GPU
def analyze_threat(threat_description, analyst_level):
    """Simple threat analysis"""
    
    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:

THREAT: {threat_description}

Provide a {analyst_level} level security analysis including:
- Threat assessment
- Potential impact
- Recommended actions

ANALYSIS:"""
    
    if pipe:
        try:
            result = pipe(
                prompt,
                max_new_tokens=300,
                do_sample=True,
                temperature=0.3,
                top_p=0.9,
                repetition_penalty=1.1
            )
            
            analysis = result[0]['generated_text'][len(prompt):].strip()
            
            if len(analysis) < 30:
                analysis = get_simple_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)}"
    else:
        analysis = get_simple_fallback(threat_description, analyst_level)
    
    processing_time = round(time.time() - start_time, 2)
    status = f"✅ Analysis completed in {processing_time}s | {model_status}"
    
    return analysis, status

def get_simple_fallback(threat_description, analyst_level):
    """Simple fallback analysis"""
    
    if analyst_level == "L1":
        return f"""🚨 L1 TRIAGE ANALYSIS

THREAT 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

PRIORITY: Immediate containment and escalation decision required"""

    elif analyst_level == "L2":
        return f"""🔍 L2 INVESTIGATION ANALYSIS

THREAT 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

CONTAINMENT:
• Implement network segmentation
• Deploy additional monitoring
• Review authentication logs
• Check for lateral movement

NEXT STEPS:
• Continue monitoring for related activity
• Update security controls as needed
• Consider L3 escalation for complex threats"""

    else:  # L3
        return f"""🎯 L3 EXPERT ANALYSIS

STRATEGIC THREAT ASSESSMENT:
{threat_description}

ADVANCED ANALYSIS:
• Threat actor attribution assessment
• Campaign analysis and TTPs
• Business impact evaluation
• Risk quantification

STRATEGIC RESPONSE:
• Coordinate incident response team
• Executive briefing preparation
• Regulatory compliance review
• Long-term security posture improvements

RECOMMENDATIONS:
• Implement advanced threat hunting
• Enhance detection capabilities
• Review security architecture
• Consider external forensics support"""

# Create simple interface
with gr.Blocks(title="Simple SOC Analyzer", theme=gr.themes.Soft(), css=simple_css) as demo:
    
    # Simple header
    gr.Markdown("""
    # 🛡️ SOC Threat Analyzer
    **Simple • Fast • Effective**
    
    Enter any security threat and get instant AI analysis.
    """)
    
    # Model status
    status_display = gr.Textbox(
        value="🔄 Loading model...",
        label="System Status",
        interactive=False,
        elem_classes=["status-box"]
    )
    
    # 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
            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
            """)
        
        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..."
            )
            
            # 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")
    
    # Simple footer
    gr.Markdown("""
    ---
    **💡 Tips:** Be specific about what you observed, include timestamps, IP addresses, user accounts, or file names when available.
    """)
    
    # Event handlers
    analyze_btn.click(
        fn=analyze_threat,
        inputs=[threat_input, analyst_level],
        outputs=[analysis_output, process_status]
    )
    
    # Initialize model on startup
    demo.load(
        fn=load_model,
        outputs=[status_display]
    )

if __name__ == "__main__":
    demo.launch(share=True)