#!/usr/bin/env python3
"""
Initialize SIGMA rule templates for enhanced rule generation
"""

import json
from datetime import datetime
from main import SessionLocal, RuleTemplate, Base, engine

# Create tables if they don't exist
Base.metadata.create_all(bind=engine)

# Template definitions with actual SIGMA rule content
SIGMA_TEMPLATES = [
    {
        "template_name": "Process Execution Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: process_creation
    product: windows
detection:
    selection:
        Image|endswith:
{{PROCESSES}}
    selection_cmd:
        CommandLine|contains:
{{COMMANDS}}
    condition: selection or selection_cmd
falsepositives:
    - Legitimate software installations
    - System administration tasks
level: {{LEVEL}}""",
        "applicable_product_patterns": ["windows", "microsoft", "office", "exchange", "sharepoint"],
        "description": "Detects suspicious process execution based on PoC exploit indicators"
    },
    {
        "template_name": "Network Connection Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: network_connection
    product: windows
detection:
    selection:
        Initiated: true
        DestinationIp:
{{NETWORK}}
    selection_url:
        DestinationHostname|contains:
{{URLS}}
    condition: selection or selection_url
falsepositives:
    - Legitimate network connections
    - Software updates
level: {{LEVEL}}""",
        "applicable_product_patterns": ["network", "web", "http", "https", "tcp", "udp"],
        "description": "Detects suspicious network connections based on PoC exploit indicators"
    },
    {
        "template_name": "File System Activity Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: file_event
    product: windows
detection:
    selection:
        TargetFilename|contains:
{{FILES}}
    condition: selection
falsepositives:
    - Legitimate file operations
    - Software installations
level: {{LEVEL}}""",
        "applicable_product_patterns": ["file", "filesystem", "upload", "download"],
        "description": "Detects suspicious file system activity based on PoC exploit indicators"
    },
    {
        "template_name": "PowerShell Execution Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: process_creation
    product: windows
detection:
    selection:
        Image|endswith:
            - '\\powershell.exe'
            - '\\pwsh.exe'
        CommandLine|contains:
{{COMMANDS}}
    condition: selection
falsepositives:
    - Legitimate PowerShell scripts
    - System administration
level: {{LEVEL}}""",
        "applicable_product_patterns": ["powershell", "windows", "microsoft"],
        "description": "Detects suspicious PowerShell execution based on PoC exploit indicators"
    },
    {
        "template_name": "Web Application Attack Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: webserver
detection:
    selection:
        cs-uri-query|contains:
{{URLS}}
    selection_user_agent:
        cs-user-agent|contains:
{{COMMANDS}}
    condition: selection or selection_user_agent
falsepositives:
    - Legitimate web application usage
    - Security scanners
level: {{LEVEL}}""",
        "applicable_product_patterns": ["web", "http", "apache", "nginx", "iis"],
        "description": "Detects web application attacks based on PoC exploit indicators"
    },
    {
        "template_name": "Registry Modification Detection",
        "template_content": """title: {{TITLE}}
id: {{RULE_ID}}
status: experimental
description: {{DESCRIPTION}}
author: CVE-SIGMA Auto Generator
date: {{DATE}}
references:
{{REFERENCES}}
tags:
{{TAGS}}
logsource:
    category: registry_event
    product: windows
detection:
    selection:
        TargetObject|contains:
{{REGISTRY}}
    condition: selection
falsepositives:
    - Legitimate software configuration changes
    - System updates
level: {{LEVEL}}""",
        "applicable_product_patterns": ["registry", "windows", "microsoft"],
        "description": "Detects suspicious registry modifications based on PoC exploit indicators"
    }
]

def initialize_templates():
    """Initialize rule templates in the database"""
    db = SessionLocal()
    
    try:
        # Clear existing templates
        db.query(RuleTemplate).delete()
        
        # Add new templates
        for template_data in SIGMA_TEMPLATES:
            template = RuleTemplate(
                template_name=template_data["template_name"],
                template_content=template_data["template_content"],
                applicable_product_patterns=template_data["applicable_product_patterns"],
                description=template_data["description"]
            )
            db.add(template)
        
        db.commit()
        print(f"Successfully initialized {len(SIGMA_TEMPLATES)} rule templates")
        
    except Exception as e:
        db.rollback()
        print(f"Error initializing templates: {e}")
        raise
    finally:
        db.close()

if __name__ == "__main__":
    initialize_templates()