bpmcdevitt
a6fb367ed4
- Extract database models from monolithic main.py (2,373 lines) into organized modules - Implement service layer pattern with dedicated business logic classes - Split API endpoints into modular FastAPI routers by functionality - Add centralized configuration management with environment variable handling - Create proper separation of concerns across data, service, and presentation layers **Architecture Changes:** - models/: SQLAlchemy database models (CVE, SigmaRule, RuleTemplate, BulkProcessingJob) - config/: Centralized settings and database configuration - services/: Business logic (CVEService, SigmaRuleService, GitHubExploitAnalyzer) - routers/: Modular API endpoints (cves, sigma_rules, bulk_operations, llm_operations) - schemas/: Pydantic request/response models **Key Improvements:** - 95% reduction in main.py size (2,373 → 120 lines) - Updated 15+ backend files with proper import structure - Eliminated circular dependencies and tight coupling - Enhanced testability with isolated service components - Better code organization for team collaboration **Backward Compatibility:** - All API endpoints maintain same URLs and behavior - Zero breaking changes to existing functionality - Database schema unchanged - Environment variables preserved 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
211 lines
No EOL
9.3 KiB
Python
211 lines
No EOL
9.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for enhanced SIGMA rule generation
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from datetime import datetime
|
|
from config.database 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()) |