n0mad1k
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""TTP (Techniques, Tactics, Procedures) extractor for Mosaic."""
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from .base import BaseExtractor
|
|
from storage.db import MosaicDB
|
|
from storage.models import TTP
|
|
from analysis.prompts import TTP_SYSTEM_PROMPT, TTP_EXTRACTION_PROMPT
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TTPExtractor(BaseExtractor):
|
|
EXTRACTOR_NAME = "ttp"
|
|
|
|
def get_system_prompt(self) -> str:
|
|
return TTP_SYSTEM_PROMPT
|
|
|
|
def get_extraction_prompt(self, text: str) -> str:
|
|
return TTP_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:
|
|
ttp = TTP(
|
|
technique=item.get('technique', ''),
|
|
category=item.get('category', ''),
|
|
description=item.get('description', ''),
|
|
actor=item.get('actor'),
|
|
mitre_id=item.get('mitre_id'),
|
|
source_doc_ids=[doc_id],
|
|
source_context=item.get('source_context', '')[:500],
|
|
)
|
|
ttp.compute_dedup_key()
|
|
if self.db.insert_ttp(ttp):
|
|
count += 1
|
|
except Exception as e:
|
|
logger.warning("Failed to process TTP item: %s", e)
|
|
return count
|