- 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>
164 lines
No EOL
5.8 KiB
Python
164 lines
No EOL
5.8 KiB
Python
from typing import List
|
|
from fastapi import APIRouter, HTTPException, Depends, BackgroundTasks
|
|
from sqlalchemy.orm import Session
|
|
from datetime import datetime, timedelta
|
|
|
|
from config.database import get_db
|
|
from models import CVE
|
|
from schemas import CVEResponse
|
|
from services import CVEService, SigmaRuleService
|
|
|
|
router = APIRouter(prefix="/api", tags=["cves"])
|
|
|
|
|
|
@router.get("/cves", response_model=List[CVEResponse])
|
|
async def get_cves(skip: int = 0, limit: int = 50, db: Session = Depends(get_db)):
|
|
"""Get all CVEs with pagination"""
|
|
cve_service = CVEService(db)
|
|
cves = cve_service.get_all_cves(limit=limit, offset=skip)
|
|
|
|
# Convert to response format
|
|
result = []
|
|
for cve in cves:
|
|
cve_dict = {
|
|
'id': str(cve.id),
|
|
'cve_id': cve.cve_id,
|
|
'description': cve.description,
|
|
'cvss_score': float(cve.cvss_score) if cve.cvss_score else None,
|
|
'severity': cve.severity,
|
|
'published_date': cve.published_date,
|
|
'affected_products': cve.affected_products,
|
|
'reference_urls': cve.reference_urls
|
|
}
|
|
result.append(CVEResponse(**cve_dict))
|
|
return result
|
|
|
|
|
|
@router.get("/cves/{cve_id}", response_model=CVEResponse)
|
|
async def get_cve(cve_id: str, db: Session = Depends(get_db)):
|
|
"""Get specific CVE by ID"""
|
|
cve_service = CVEService(db)
|
|
cve = cve_service.get_cve_by_id(cve_id)
|
|
|
|
if not cve:
|
|
raise HTTPException(status_code=404, detail="CVE not found")
|
|
|
|
cve_dict = {
|
|
'id': str(cve.id),
|
|
'cve_id': cve.cve_id,
|
|
'description': cve.description,
|
|
'cvss_score': float(cve.cvss_score) if cve.cvss_score else None,
|
|
'severity': cve.severity,
|
|
'published_date': cve.published_date,
|
|
'affected_products': cve.affected_products,
|
|
'reference_urls': cve.reference_urls
|
|
}
|
|
return CVEResponse(**cve_dict)
|
|
|
|
|
|
@router.post("/fetch-cves")
|
|
async def manual_fetch_cves(background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
|
|
"""Manually trigger CVE fetch from NVD API"""
|
|
async def fetch_task():
|
|
try:
|
|
cve_service = CVEService(db)
|
|
sigma_service = SigmaRuleService(db)
|
|
|
|
print("Manual CVE fetch initiated...")
|
|
# Use 30 days for manual fetch to get more results
|
|
new_cves = await cve_service.fetch_recent_cves(days_back=30)
|
|
|
|
rules_generated = 0
|
|
for cve in new_cves:
|
|
sigma_rule = sigma_service.generate_sigma_rule(cve)
|
|
if sigma_rule:
|
|
rules_generated += 1
|
|
|
|
db.commit()
|
|
print(f"Manual fetch complete: {len(new_cves)} CVEs, {rules_generated} rules generated")
|
|
except Exception as e:
|
|
print(f"Manual fetch error: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
background_tasks.add_task(fetch_task)
|
|
return {"message": "CVE fetch initiated (30-day lookback)", "status": "started"}
|
|
|
|
|
|
@router.get("/cve-stats")
|
|
async def get_cve_stats(db: Session = Depends(get_db)):
|
|
"""Get CVE statistics"""
|
|
cve_service = CVEService(db)
|
|
return cve_service.get_cve_stats()
|
|
|
|
|
|
@router.get("/test-nvd")
|
|
async def test_nvd_connection():
|
|
"""Test endpoint to check NVD API connectivity"""
|
|
try:
|
|
import requests
|
|
import os
|
|
|
|
# Test with a simple request using current date
|
|
end_date = datetime.utcnow()
|
|
start_date = end_date - timedelta(days=30)
|
|
|
|
url = "https://services.nvd.nist.gov/rest/json/cves/2.0/"
|
|
params = {
|
|
"lastModStartDate": start_date.strftime("%Y-%m-%dT%H:%M:%S.000+00:00"),
|
|
"lastModEndDate": end_date.strftime("%Y-%m-%dT%H:%M:%S.000+00:00"),
|
|
"resultsPerPage": 5,
|
|
"startIndex": 0
|
|
}
|
|
|
|
headers = {
|
|
"User-Agent": "CVE-SIGMA-Generator/1.0",
|
|
"Accept": "application/json"
|
|
}
|
|
|
|
nvd_api_key = os.getenv("NVD_API_KEY")
|
|
if nvd_api_key:
|
|
headers["apiKey"] = nvd_api_key
|
|
|
|
print(f"Testing NVD API with URL: {url}")
|
|
print(f"Test params: {params}")
|
|
|
|
response = requests.get(url, params=params, headers=headers, timeout=30)
|
|
|
|
print(f"Response status: {response.status_code}")
|
|
print(f"Response headers: {dict(response.headers)}")
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
total_results = data.get("totalResults", 0)
|
|
vulnerabilities = data.get("vulnerabilities", [])
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": f"Successfully connected to NVD API. Found {total_results} total results, returned {len(vulnerabilities)} vulnerabilities.",
|
|
"total_results": total_results,
|
|
"returned_count": len(vulnerabilities),
|
|
"has_api_key": bool(nvd_api_key),
|
|
"rate_limit": "50 requests/30s" if nvd_api_key else "5 requests/30s"
|
|
}
|
|
else:
|
|
response_text = response.text[:500] # Limit response text
|
|
return {
|
|
"status": "error",
|
|
"message": f"NVD API returned status {response.status_code}",
|
|
"response_preview": response_text,
|
|
"has_api_key": bool(nvd_api_key)
|
|
}
|
|
|
|
except requests.RequestException as e:
|
|
return {
|
|
"status": "error",
|
|
"message": f"Network error connecting to NVD API: {str(e)}",
|
|
"has_api_key": bool(os.getenv("NVD_API_KEY"))
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"status": "error",
|
|
"message": f"Unexpected error: {str(e)}",
|
|
"has_api_key": bool(os.getenv("NVD_API_KEY"))
|
|
} |