#!/usr/bin/env python3
"""MCP tool definitions for CVE database operations."""
import json
from pathlib import Path
import sys
# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
from database.db import init_db, get_cve, search_cves, get_stats
# Initialize database connection
_conn = None
def get_connection():
"""Get or create database connection."""
global _conn
if _conn is None:
_conn = init_db()
return _conn
def tool_get_cve_details(cve_id: str) -> str:
"""
Get detailed information about a specific CVE.
Args:
cve_id: The CVE identifier (e.g., "CVE-2024-0001")
Returns:
JSON string with CVE details or error message
"""
conn = get_connection()
# Normalize CVE ID format
cve_id = cve_id.strip().upper()
if not cve_id.startswith("CVE-"):
cve_id = f"CVE-{cve_id}"
result = get_cve(conn, cve_id)
if result:
# Parse references back to list
if result.get('references_json'):
try:
result['references'] = json.loads(result['references_json'])
except json.JSONDecodeError:
result['references'] = []
del result['references_json']
return json.dumps(result, indent=2)
else:
return json.dumps({
"error": "CVE not found",
"cve_id": cve_id,
"suggestion": "Try searching with search_cves tool"
})
def tool_search_cves(keyword: str, limit: int = 10) -> str:
"""
Search for CVEs by keyword in description.
Args:
keyword: Search term to look for in CVE descriptions
limit: Maximum number of results (default: 10, max: 50)
Returns:
JSON string with list of matching CVEs
"""
conn = get_connection()
# Validate limit
limit = min(max(1, limit), 50)
results = search_cves(conn, keyword.strip(), limit)
# Format results
formatted = []
for cve in results:
formatted.append({
"cve_id": cve['cve_id'],
"severity": cve['severity'],
"cvss_score": cve['cvss_score'],
"description": cve['description'][:300] + "..." if len(cve['description']) > 300 else cve['description'],
"published_date": cve['published_date']
})
return json.dumps({
"query": keyword,
"count": len(formatted),
"results": formatted
}, indent=2)
def tool_get_statistics() -> str:
"""
Get database statistics and metadata.
Returns:
JSON string with database stats
"""
conn = get_connection()
stats = get_stats(conn)
return json.dumps({
"database_info": {
"total_cves": stats['total_cves'],
"date_range": {
"oldest": stats['oldest_cve'],
"newest": stats['newest_cve']
},
"last_update": stats['last_update']
}
}, indent=2)
