import gradio as gr
import os
import json
import datetime
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import yaml
import uuid
import tempfile
import shutil

# Demo configuration
DEMO_CASE_ID = f"DEMO-{uuid.uuid4().hex[:8]}"
DEMO_OUTPUT_DIR = "demo_output"
DEMO_EVIDENCE_DIR = os.path.join(DEMO_OUTPUT_DIR, "evidence")
DEMO_ANALYSIS_DIR = os.path.join(DEMO_OUTPUT_DIR, "analysis")
DEMO_REPORT_DIR = os.path.join(DEMO_OUTPUT_DIR, "reports")

# Create directories if they don't exist
os.makedirs(DEMO_EVIDENCE_DIR, exist_ok=True)
os.makedirs(DEMO_ANALYSIS_DIR, exist_ok=True)
os.makedirs(DEMO_REPORT_DIR, exist_ok=True)

# Cloud provider connection functions
def test_aws_connection(access_key, secret_key, region):
    """Test connection to AWS"""
    try:
        import boto3
        session = boto3.Session(
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name=region
        )
        sts = session.client('sts')
        identity = sts.get_caller_identity()
        return True, f"Successfully connected to AWS as {identity['Arn']}"
    except Exception as e:
        return False, f"Failed to connect to AWS: {str(e)}"

def test_azure_connection(tenant_id, client_id, client_secret):
    """Test connection to Azure"""
    try:
        from azure.identity import ClientSecretCredential
        from azure.mgmt.resource import ResourceManagementClient
        
        credential = ClientSecretCredential(
            tenant_id=tenant_id,
            client_id=client_id,
            client_secret=client_secret
        )
        
        # Create a resource management client
        resource_client = ResourceManagementClient(credential, subscription_id)
        
        # List resource groups to test the connection
        resource_groups = list(resource_client.resource_groups.list())
        return True, f"Successfully connected to Azure. Found {len(resource_groups)} resource groups."
    except Exception as e:
        return False, f"Failed to connect to Azure: {str(e)}"

def test_gcp_connection(service_account_json):
    """Test connection to GCP"""
    try:
        import json
        from google.oauth2 import service_account
        from google.cloud import storage
        
        # Create a temporary file to store the service account JSON
        fd, path = tempfile.mkstemp()
        try:
            with os.fdopen(fd, 'w') as tmp:
                tmp.write(service_account_json)
            
            # Create credentials from the service account file
            credentials = service_account.Credentials.from_service_account_file(path)
            
            # Create a storage client to test the connection
            storage_client = storage.Client(credentials=credentials)
            
            # List buckets to test the connection
            buckets = list(storage_client.list_buckets())
            return True, f"Successfully connected to GCP. Found {len(buckets)} storage buckets."
        finally:
            os.remove(path)
    except Exception as e:
        return False, f"Failed to connect to GCP: {str(e)}"

# Sample data for demonstration
def generate_sample_data(case_info, cloud_provider, incident_type, use_real_data=False, credentials=None):
    """Generate sample data for demonstration purposes or collect real data if credentials provided"""
    
    if use_real_data and credentials:
        # This would be where we implement real data collection using the provided credentials
        # For now, we'll return a message indicating this would use real data
        return {
            "timeline": [],
            "patterns": [],
            "anomalies": [],
            "files": {},
            "message": "In a production deployment, this would collect real data from your cloud provider."
        }
    
    # Create sample timeline data
    timeline_data = []
    base_time = datetime.datetime.now() - datetime.timedelta(days=1)
    
    # Different events based on incident type
    if incident_type == "Unauthorized Access":
        events = [
            {"event": "Failed login attempt", "source": "Authentication Logs", "severity": "Low"},
            {"event": "Successful login from unusual IP", "source": "Authentication Logs", "severity": "Medium"},
            {"event": "User privilege escalation", "source": "IAM Logs", "severity": "High"},
            {"event": "Access to sensitive data", "source": "Data Access Logs", "severity": "High"},
            {"event": "Configuration change", "source": "Configuration Logs", "severity": "Medium"},
            {"event": "New API key created", "source": "IAM Logs", "severity": "High"},
            {"event": "Data download initiated", "source": "Data Access Logs", "severity": "Critical"},
            {"event": "Unusual network traffic", "source": "Network Logs", "severity": "Medium"}
        ]
    elif incident_type == "Data Exfiltration":
        events = [
            {"event": "Large query executed", "source": "Database Logs", "severity": "Medium"},
            {"event": "Unusual data access pattern", "source": "Data Access Logs", "severity": "Medium"},
            {"event": "Large data transfer initiated", "source": "Network Logs", "severity": "High"},
            {"event": "Connection to unknown external endpoint", "source": "Network Logs", "severity": "High"},
            {"event": "Storage object permissions modified", "source": "Storage Logs", "severity": "Medium"},
            {"event": "Unusual user behavior", "source": "User Activity Logs", "severity": "Medium"},
            {"event": "Data archive created", "source": "Storage Logs", "severity": "Medium"},
            {"event": "Unusual egress traffic spike", "source": "Network Logs", "severity": "Critical"}
        ]
    else:  # Ransomware
        events = [
            {"event": "Unusual process execution", "source": "System Logs", "severity": "Medium"},
            {"event": "Multiple file modifications", "source": "File System Logs", "severity": "High"},
            {"event": "Encryption library loaded", "source": "System Logs", "severity": "High"},
            {"event": "Mass file type changes", "source": "Storage Logs", "severity": "Critical"},
            {"event": "Backup deletion attempt", "source": "Backup Logs", "severity": "Critical"},
            {"event": "Unusual IAM activity", "source": "IAM Logs", "severity": "Medium"},
            {"event": "Recovery service disabled", "source": "System Logs", "severity": "High"},
            {"event": "Ransom note created", "source": "File System Logs", "severity": "Critical"}
        ]
    
    # Create timeline with timestamps
    for i, event in enumerate(events):
        event_time = base_time + datetime.timedelta(minutes=i*15)
        timeline_data.append({
            "timestamp": event_time.isoformat(),
            "event": event["event"],
            "source": event["source"],
            "cloud_provider": cloud_provider,
            "severity": event["severity"],
            "case_id": case_info["case_id"]
        })
    
    # Create patterns data
    patterns = []
    if incident_type == "Unauthorized Access":
        patterns = [
            {"pattern": "Brute Force Login Attempt", "confidence": 0.85, "matched_events": 3},
            {"pattern": "Privilege Escalation", "confidence": 0.92, "matched_events": 2}
        ]
    elif incident_type == "Data Exfiltration":
        patterns = [
            {"pattern": "Data Staging Activity", "confidence": 0.88, "matched_events": 3},
            {"pattern": "Exfiltration Over Alternative Protocol", "confidence": 0.76, "matched_events": 2}
        ]
    else:  # Ransomware
        patterns = [
            {"pattern": "Mass File Encryption", "confidence": 0.94, "matched_events": 4},
            {"pattern": "Defense Evasion", "confidence": 0.81, "matched_events": 3}
        ]
    
    # Create anomalies data
    anomalies = []
    if incident_type == "Unauthorized Access":
        anomalies = [
            {"anomaly": "Login from unusual location", "deviation": 3.6, "severity": "High"},
            {"anomaly": "Off-hours access", "deviation": 2.8, "severity": "Medium"}
        ]
    elif incident_type == "Data Exfiltration":
        anomalies = [
            {"anomaly": "Unusual data access volume", "deviation": 4.2, "severity": "High"},
            {"anomaly": "Abnormal query pattern", "deviation": 3.1, "severity": "Medium"}
        ]
    else:  # Ransomware
        anomalies = [
            {"anomaly": "Unusual file system activity", "deviation": 4.7, "severity": "Critical"},
            {"anomaly": "Suspicious process behavior", "deviation": 3.9, "severity": "High"}
        ]
    
    # Save data to files
    timeline_file = os.path.join(DEMO_EVIDENCE_DIR, f"{DEMO_CASE_ID}_timeline.json")
    patterns_file = os.path.join(DEMO_ANALYSIS_DIR, f"{DEMO_CASE_ID}_patterns.json")
    anomalies_file = os.path.join(DEMO_ANALYSIS_DIR, f"{DEMO_CASE_ID}_anomalies.json")
    
    with open(timeline_file, 'w') as f:
        json.dump(timeline_data, f, indent=2)
    
    with open(patterns_file, 'w') as f:
        json.dump(patterns, f, indent=2)
    
    with open(anomalies_file, 'w') as f:
        json.dump(anomalies, f, indent=2)
    
    return {
        "timeline": timeline_data,
        "patterns": patterns,
        "anomalies": anomalies,
        "files": {
            "timeline": timeline_file,
            "patterns": patterns_file,
            "anomalies": anomalies_file
        }
    }

def analyze_evidence(data):
    """Perform analysis on the evidence data"""
    
    # If there's no timeline data, return empty results
    if not data["timeline"]:
        return {
            "severity_counts": {},
            "source_counts": {},
            "charts": {
                "analysis": None,
                "timeline": None
            }
        }
    
    # Convert timeline to DataFrame for analysis
    timeline_df = pd.DataFrame(data["timeline"])
    timeline_df["timestamp"] = pd.to_datetime(timeline_df["timestamp"])
    
    # Sort by timestamp
    timeline_df = timeline_df.sort_values("timestamp")
    
    # Count events by severity
    severity_counts = timeline_df["severity"].value_counts()
    
    # Count events by source
    source_counts = timeline_df["source"].value_counts()
    
    # Create visualizations
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # Severity pie chart
    ax1.pie(severity_counts, labels=severity_counts.index, autopct='%1.1f%%', 
            colors=sns.color_palette("YlOrRd", len(severity_counts)))
    ax1.set_title("Events by Severity")
    
    # Source bar chart
    sns.barplot(x=source_counts.values, y=source_counts.index, ax=ax2, palette="viridis")
    ax2.set_title("Events by Source")
    ax2.set_xlabel("Count")
    
    # Save the figure
    chart_file = os.path.join(DEMO_ANALYSIS_DIR, f"{DEMO_CASE_ID}_analysis_charts.png")
    plt.tight_layout()
    plt.savefig(chart_file)
    plt.close()
    
    # Create a timeline visualization
    plt.figure(figsize=(12, 6))
    
    # Create a categorical y-axis based on source
    sources = timeline_df["source"].unique()
    source_map = {source: i for i, source in enumerate(sources)}
    timeline_df["source_num"] = timeline_df["source"].map(source_map)
    
    # Map severity to color
    severity_colors = {
        "Low": "green",
        "Medium": "blue",
        "High": "orange",
        "Critical": "red"
    }
    colors = timeline_df["severity"].map(severity_colors)
    
    # Plot the timeline
    plt.scatter(timeline_df["timestamp"], timeline_df["source_num"], c=colors, s=100)
    
    # Add event labels
    for i, row in timeline_df.iterrows():
        plt.text(row["timestamp"], row["source_num"], row["event"], 
                 fontsize=8, ha="right", va="bottom", rotation=25)
    
    plt.yticks(range(len(sources)), sources)
    plt.xlabel("Time")
    plt.ylabel("Event Source")
    plt.title("Incident Timeline")
    
    # Save the timeline
    timeline_chart = os.path.join(DEMO_ANALYSIS_DIR, f"{DEMO_CASE_ID}_timeline_chart.png")
    plt.tight_layout()
    plt.savefig(timeline_chart)
    plt.close()
    
    return {
        "severity_counts": severity_counts.to_dict(),
        "source_counts": source_counts.to_dict(),
        "charts": {
            "analysis": chart_file,
            "timeline": timeline_chart
        }
    }

def generate_report(case_info, data, analysis, report_format):
    """Generate a report based on the analysis"""
    
    # Create report content
    report = {
        "case_information": case_info,
        "executive_summary": f"This report presents the findings of a forensic investigation into a {case_info['incident_type']} incident in {case_info['cloud_provider']} cloud environment.",
        "timeline": data["timeline"],
        "patterns_detected": data["patterns"],
        "anomalies_detected": data["anomalies"],
        "analysis_results": {
            "severity_distribution": analysis.get("severity_counts", {}),
            "source_distribution": analysis.get("source_counts", {})
        },
        "recommendations": [
            "Implement multi-factor authentication for all privileged accounts",
            "Review and restrict IAM permissions following principle of least privilege",
            "Enable comprehensive logging across all cloud services",
            "Implement automated alerting for suspicious activities",
            "Conduct regular security assessments of cloud environments"
        ]
    }
    
    # Save report in requested format
    if report_format == "JSON":
        report_file = os.path.join(DEMO_REPORT_DIR, f"{DEMO_CASE_ID}_report.json")
        with open(report_file, 'w') as f:
            json.dump(report, f, indent=2)
    else:  # HTML
        # Create a simple HTML report
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Forensic Analysis Report - {case_info['case_id']}</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 40px; }}
                h1, h2, h3 {{ color: #2c3e50; }}
                .section {{ margin-bottom: 30px; }}
                .severity-high {{ color: #e74c3c; }}
                .severity-medium {{ color: #f39c12; }}
                .severity-low {{ color: #27ae60; }}
                table {{ border-collapse: collapse; width: 100%; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
                tr:nth-child(even) {{ background-color: #f9f9f9; }}
                .chart-container {{ display: flex; justify-content: center; margin: 20px 0; }}
                .chart {{ max-width: 100%; height: auto; margin: 10px; }}
                .message {{ background-color: #f8f9fa; padding: 15px; border-left: 5px solid #4e73df; margin-bottom: 20px; }}
            </style>
        </head>
        <body>
            <h1>Cloud Forensics Analysis Report</h1>
            
            <div class="section">
                <h2>Case Information</h2>
                <p><strong>Case ID:</strong> {case_info['case_id']}</p>
                <p><strong>Investigator:</strong> {case_info['investigator']}</p>
                <p><strong>Organization:</strong> {case_info['organization']}</p>
                <p><strong>Cloud Provider:</strong> {case_info['cloud_provider']}</p>
                <p><strong>Incident Type:</strong> {case_info['incident_type']}</p>
                <p><strong>Report Date:</strong> {datetime.datetime.now().strftime('%Y-%m-%d')}</p>
            </div>
            
            <div class="section">
                <h2>Executive Summary</h2>
                <p>{report['executive_summary']}</p>
        """
        
        # Add message if using real data
        if "message" in data:
            html_content += f"""
                <div class="mes
(Content truncated due to size limit. Use line ranges to read in chunks)