n0mad1k
52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
"""Infrastructure indicator extractor for Mosaic (IPs, domains, C2, hashes)."""
|
|
import logging
|
|
|
|
from .base import BaseExtractor
|
|
from storage.models import Infrastructure
|
|
from analysis.prompts import INFRASTRUCTURE_SYSTEM_PROMPT, INFRASTRUCTURE_EXTRACTION_PROMPT
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class InfrastructureExtractor(BaseExtractor):
|
|
EXTRACTOR_NAME = "infrastructure"
|
|
|
|
def __init__(self, *args, collection: str = "", doc_title: str = "",
|
|
doc_date: str = None, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.collection = collection
|
|
self.doc_title = doc_title
|
|
self.doc_date = doc_date
|
|
|
|
def get_system_prompt(self) -> str:
|
|
return INFRASTRUCTURE_SYSTEM_PROMPT
|
|
|
|
def get_extraction_prompt(self, text: str) -> str:
|
|
return INFRASTRUCTURE_EXTRACTION_PROMPT.format(document=text)
|
|
|
|
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:
|
|
infra = Infrastructure(
|
|
indicator=item.get('indicator', ''),
|
|
indicator_type=item.get('indicator_type', ''),
|
|
context=item.get('context', ''),
|
|
source_doc_ids=[doc_id],
|
|
source_context=item.get('source_context', '')[:500],
|
|
collection=self.collection,
|
|
doc_title=self.doc_title,
|
|
doc_date=self.doc_date,
|
|
)
|
|
self.db.insert_infrastructure(infra)
|
|
count += 1
|
|
except Exception as e:
|
|
logger.warning("Failed to process infrastructure item: %s", e)
|
|
return count
|