n0mad1k
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
"""Named tool/exploit/implant extractor for Mosaic."""
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from .base import BaseExtractor
|
|
from storage.db import MosaicDB
|
|
from storage.models import ExtractedTool
|
|
from analysis.prompts import TOOL_SYSTEM_PROMPT, TOOL_EXTRACTION_PROMPT
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
KNOWN_TOOLS_PATH = Path(__file__).parent.parent / 'data' / 'known_tools.json'
|
|
|
|
|
|
class ToolExtractor(BaseExtractor):
|
|
EXTRACTOR_NAME = "tool"
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.known_tools = self._load_known_tools()
|
|
|
|
def _load_known_tools(self) -> list[str]:
|
|
try:
|
|
with open(KNOWN_TOOLS_PATH) as f:
|
|
data = json.load(f)
|
|
names = []
|
|
for tool in data.get('tools', []):
|
|
names.append(tool['name'].lower())
|
|
for alias in tool.get('aliases', []):
|
|
names.append(alias.lower())
|
|
return names
|
|
except Exception as e:
|
|
logger.warning("Could not load known tools: %s", e)
|
|
return []
|
|
|
|
def get_system_prompt(self) -> str:
|
|
return TOOL_SYSTEM_PROMPT
|
|
|
|
def get_extraction_prompt(self, text: str) -> str:
|
|
# Include known tools as context
|
|
known_hint = ', '.join(self.known_tools[:50]) if self.known_tools else 'none loaded'
|
|
return TOOL_EXTRACTION_PROMPT.format(document=text, known_tools=known_hint)
|
|
|
|
def get_json_schema(self) -> dict:
|
|
return {
|
|
"properties": {
|
|
"items": {"type": "array"},
|
|
},
|
|
"required": ["items"]
|
|
}
|
|
|
|
def process_response(self, data: dict, doc_id: str) -> int:
|
|
items = data.get('items', [])
|
|
if isinstance(data, list):
|
|
items = data
|
|
count = 0
|
|
for item in items:
|
|
try:
|
|
tool = ExtractedTool(
|
|
name=item.get('name', ''),
|
|
aliases=item.get('aliases', []),
|
|
description=item.get('description', ''),
|
|
capability=item.get('capability', ''),
|
|
target_platforms=item.get('target_platforms', []),
|
|
target_software=item.get('target_software', []),
|
|
cves=item.get('cves', []),
|
|
source_doc_ids=[doc_id],
|
|
source_context=item.get('source_context', '')[:500],
|
|
mitre_techniques=item.get('mitre_techniques', []),
|
|
)
|
|
tool.compute_dedup_key()
|
|
if self.db.insert_extracted_tool(tool):
|
|
count += 1
|
|
except Exception as e:
|
|
logger.warning("Failed to process tool item: %s", e)
|
|
return count
|