#!/usr/bin/env python3 """ Test script for enhanced SIGMA rule generation """ import asyncio import json from datetime import datetime from main import SessionLocal, CVE, SigmaRule, Base, engine from enhanced_sigma_generator import EnhancedSigmaGenerator from nomi_sec_client import NomiSecClient from initialize_templates import initialize_templates # Create tables if they don't exist Base.metadata.create_all(bind=engine) async def test_enhanced_rule_generation(): """Test the enhanced rule generation with mock data""" # Initialize templates print("Initializing templates...") initialize_templates() db = SessionLocal() try: # Check if CVE already exists, if not create it test_cve = db.query(CVE).filter(CVE.cve_id == "CVE-2014-7236").first() if not test_cve: # Create a test CVE with mock PoC data test_cve = CVE( cve_id="CVE-2014-7236", description="Remote code execution vulnerability in Microsoft Office", cvss_score=8.5, severity="high", published_date=datetime(2014, 10, 15), affected_products=["Microsoft Office", "Windows"], poc_count=2, poc_data=[ { "id": "test1", "name": "CVE-2014-7236-exploit", "owner": "security-researcher", "full_name": "security-researcher/CVE-2014-7236-exploit", "html_url": "https://github.com/security-researcher/CVE-2014-7236-exploit", "description": "PowerShell exploit for CVE-2014-7236 using cmd.exe and powershell.exe", "stargazers_count": 15, "created_at": "2014-11-01T00:00:00Z", "updated_at": "2014-11-15T00:00:00Z", "quality_analysis": { "quality_score": 75, "quality_tier": "good", "factors": { "star_score": 30, "recency_score": 10, "description_score": 15, "vuln_description_score": 15, "name_relevance_score": 10 } }, "exploit_indicators": { "processes": ["powershell.exe", "cmd.exe"], "files": ["exploit.ps1", "payload.exe"], "commands": ["Invoke-Expression", "DownloadString", "whoami"], "network": ["192.168.1.100", "8080"], "urls": ["http://malicious.com/payload"], "registry": ["HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft"] } }, { "id": "test2", "name": "office-exploit-poc", "owner": "hacker", "full_name": "hacker/office-exploit-poc", "html_url": "https://github.com/hacker/office-exploit-poc", "description": "Office document exploit with malicious macro", "stargazers_count": 8, "created_at": "2014-12-01T00:00:00Z", "updated_at": "2014-12-10T00:00:00Z", "quality_analysis": { "quality_score": 45, "quality_tier": "fair", "factors": { "star_score": 16, "recency_score": 8, "description_score": 12, "vuln_description_score": 0, "name_relevance_score": 5 } }, "exploit_indicators": { "processes": ["winword.exe", "excel.exe"], "files": ["document.docx", "malicious.xlsm"], "commands": ["CreateObject", "Shell.Application"], "network": ["10.0.0.1"], "urls": ["http://evil.com/download"], "registry": ["HKEY_CURRENT_USER\\Software\\Microsoft\\Office"] } } ] ) # Add to database db.add(test_cve) db.commit() else: # Update existing CVE with our mock PoC data test_cve.poc_count = 2 test_cve.poc_data = [ { "id": "test1", "name": "CVE-2014-7236-exploit", "owner": "security-researcher", "full_name": "security-researcher/CVE-2014-7236-exploit", "html_url": "https://github.com/security-researcher/CVE-2014-7236-exploit", "description": "PowerShell exploit for CVE-2014-7236 using cmd.exe and powershell.exe", "stargazers_count": 15, "created_at": "2014-11-01T00:00:00Z", "updated_at": "2014-11-15T00:00:00Z", "quality_analysis": { "quality_score": 75, "quality_tier": "good", "factors": { "star_score": 30, "recency_score": 10, "description_score": 15, "vuln_description_score": 15, "name_relevance_score": 10 } }, "exploit_indicators": { "processes": ["powershell.exe", "cmd.exe"], "files": ["exploit.ps1", "payload.exe"], "commands": ["Invoke-Expression", "DownloadString", "whoami"], "network": ["192.168.1.100", "8080"], "urls": ["http://malicious.com/payload"], "registry": ["HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft"] } }, { "id": "test2", "name": "office-exploit-poc", "owner": "hacker", "full_name": "hacker/office-exploit-poc", "html_url": "https://github.com/hacker/office-exploit-poc", "description": "Office document exploit with malicious macro", "stargazers_count": 8, "created_at": "2014-12-01T00:00:00Z", "updated_at": "2014-12-10T00:00:00Z", "quality_analysis": { "quality_score": 45, "quality_tier": "fair", "factors": { "star_score": 16, "recency_score": 8, "description_score": 12, "vuln_description_score": 0, "name_relevance_score": 5 } }, "exploit_indicators": { "processes": ["winword.exe", "excel.exe"], "files": ["document.docx", "malicious.xlsm"], "commands": ["CreateObject", "Shell.Application"], "network": ["10.0.0.1"], "urls": ["http://evil.com/download"], "registry": ["HKEY_CURRENT_USER\\Software\\Microsoft\\Office"] } } ] db.commit() print(f"Using CVE: {test_cve.cve_id} with {test_cve.poc_count} PoCs") # Generate enhanced rule print("Generating enhanced SIGMA rule...") generator = EnhancedSigmaGenerator(db) result = await generator.generate_enhanced_rule(test_cve) print(f"Generation result: {result}") if result.get('success'): # Fetch the generated rule sigma_rule = db.query(SigmaRule).filter(SigmaRule.cve_id == test_cve.cve_id).first() if sigma_rule: print("\n" + "="*60) print("GENERATED SIGMA RULE:") print("="*60) print(sigma_rule.rule_content) print("="*60) print(f"Detection Type: {sigma_rule.detection_type}") print(f"Log Source: {sigma_rule.log_source}") print(f"Confidence Level: {sigma_rule.confidence_level}") print(f"PoC Quality Score: {sigma_rule.poc_quality_score}") print(f"Exploit Indicators: {sigma_rule.exploit_indicators}") print("="*60) else: print("No SIGMA rule found in database") else: print(f"Rule generation failed: {result.get('error')}") except Exception as e: print(f"Error during test: {e}") import traceback traceback.print_exc() finally: db.close() if __name__ == "__main__": asyncio.run(test_enhanced_rule_generation())