95 lines
No EOL
2.9 KiB
Python
95 lines
No EOL
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Initialize SIGMA rule templates for enhanced rule generation
|
|
"""
|
|
|
|
import json
|
|
import yaml
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from main import SessionLocal, RuleTemplate, Base, engine
|
|
|
|
# Create tables if they don't exist
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
def load_templates_from_files():
|
|
"""Load SIGMA rule templates from YAML files in the templates directory"""
|
|
templates = []
|
|
templates_dir = Path(__file__).parent / "templates"
|
|
|
|
if not templates_dir.exists():
|
|
print(f"Templates directory not found: {templates_dir}")
|
|
return templates
|
|
|
|
# Load all YAML files from templates directory
|
|
for template_file in templates_dir.glob("*.yaml"):
|
|
try:
|
|
with open(template_file, 'r', encoding='utf-8') as f:
|
|
template_data = yaml.safe_load(f)
|
|
templates.append(template_data)
|
|
print(f"Loaded template: {template_data['template_name']}")
|
|
except Exception as e:
|
|
print(f"Error loading template from {template_file}: {e}")
|
|
|
|
return templates
|
|
|
|
def initialize_templates():
|
|
"""Initialize rule templates in the database"""
|
|
db = SessionLocal()
|
|
|
|
try:
|
|
# Load templates from YAML files
|
|
templates = load_templates_from_files()
|
|
|
|
if not templates:
|
|
print("No templates found to initialize")
|
|
return
|
|
|
|
# Clear existing templates
|
|
db.query(RuleTemplate).delete()
|
|
|
|
# Add new templates
|
|
for template_data in 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(templates)} rule templates")
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
print(f"Error initializing templates: {e}")
|
|
raise
|
|
finally:
|
|
db.close()
|
|
|
|
def list_available_templates():
|
|
"""List all available template files and their details"""
|
|
templates = load_templates_from_files()
|
|
|
|
if not templates:
|
|
print("No templates found")
|
|
return
|
|
|
|
print(f"\nFound {len(templates)} available templates:")
|
|
print("=" * 60)
|
|
|
|
for i, template in enumerate(templates, 1):
|
|
print(f"{i}. {template['template_name']}")
|
|
print(f" Description: {template['description']}")
|
|
print(f" Applicable patterns: {', '.join(template['applicable_product_patterns'])}")
|
|
print()
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) > 1 and sys.argv[1] == "list":
|
|
list_available_templates()
|
|
else:
|
|
initialize_templates() |