mirror of
https://github.com/overcuriousity/trace.git
synced 2025-12-20 04:52:21 +00:00
b6387f4b0c
Major refactoring to organize code into focused, single-responsibility modules that are easier for AI coding agents and developers to navigate and modify. **Module Reorganization:** Models Package (trace/models/): - Moved models.py content into models/__init__.py - Extracted IOC extraction into models/extractors/ioc_extractor.py (236 lines) - Extracted tag extraction into models/extractors/tag_extractor.py (34 lines) - Reduced duplication and improved maintainability Storage Package (trace/storage_impl/): - Split storage.py (402 lines) into focused modules: - storage.py: Main Storage class (112 lines) - state_manager.py: StateManager for context/settings (92 lines) - lock_manager.py: Cross-platform file locking (87 lines) - demo_data.py: Demo case creation (143 lines) - Added backward-compatible wrapper at trace/storage.py TUI Utilities (trace/tui/): - Created rendering package: - colors.py: Color pair constants and initialization (43 lines) - text_renderer.py: Text rendering with highlighting (137 lines) - Created handlers package: - export_handler.py: Export functionality (238 lines) - Main tui.py (3307 lines) remains for future refactoring **Benefits:** - Smaller, focused files (most < 250 lines) - Clear single responsibilities - Easier to locate and modify specific functionality - Better separation of concerns - Reduced cognitive load for AI agents - All tests pass, no features removed **Testing:** - All existing tests pass - Imports verified - CLI and storage functionality tested - Backward compatibility maintained Updated CLAUDE.md to document new architecture and AI optimization strategy.
132 lines
4.1 KiB
Python
132 lines
4.1 KiB
Python
"""Data models for trace application"""
|
|
|
|
import time
|
|
import hashlib
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Dict
|
|
|
|
from .extractors import TagExtractor, IOCExtractor
|
|
|
|
|
|
@dataclass
|
|
class Note:
|
|
content: str
|
|
timestamp: float = field(default_factory=time.time)
|
|
note_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
|
content_hash: str = ""
|
|
signature: Optional[str] = None
|
|
tags: List[str] = field(default_factory=list)
|
|
iocs: List[str] = field(default_factory=list)
|
|
|
|
def extract_tags(self):
|
|
"""Extract hashtags from content (case-insensitive, stored lowercase)"""
|
|
self.tags = TagExtractor.extract_tags(self.content)
|
|
|
|
def extract_iocs(self):
|
|
"""Extract Indicators of Compromise from content"""
|
|
self.iocs = IOCExtractor.extract_iocs(self.content)
|
|
|
|
def calculate_hash(self):
|
|
# We hash the content + timestamp to ensure integrity of 'when' it was said
|
|
data = f"{self.timestamp}:{self.content}".encode('utf-8')
|
|
self.content_hash = hashlib.sha256(data).hexdigest()
|
|
|
|
@staticmethod
|
|
def extract_iocs_from_text(text):
|
|
"""Extract IOCs from text and return as list of (ioc, type) tuples"""
|
|
return IOCExtractor.extract_iocs_with_types(text)
|
|
|
|
@staticmethod
|
|
def extract_iocs_with_positions(text):
|
|
"""Extract IOCs with their positions for highlighting. Returns list of (text, start, end, type) tuples"""
|
|
return IOCExtractor.extract_iocs_with_positions(text)
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"note_id": self.note_id,
|
|
"content": self.content,
|
|
"timestamp": self.timestamp,
|
|
"content_hash": self.content_hash,
|
|
"signature": self.signature,
|
|
"tags": self.tags,
|
|
"iocs": self.iocs
|
|
}
|
|
|
|
@staticmethod
|
|
def from_dict(data):
|
|
note = Note(
|
|
content=data["content"],
|
|
timestamp=data["timestamp"],
|
|
note_id=data["note_id"],
|
|
content_hash=data.get("content_hash", ""),
|
|
signature=data.get("signature"),
|
|
tags=data.get("tags", []),
|
|
iocs=data.get("iocs", [])
|
|
)
|
|
return note
|
|
|
|
|
|
@dataclass
|
|
class Evidence:
|
|
name: str
|
|
evidence_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
|
description: str = ""
|
|
metadata: Dict[str, str] = field(default_factory=dict)
|
|
notes: List[Note] = field(default_factory=list)
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"evidence_id": self.evidence_id,
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"metadata": self.metadata,
|
|
"notes": [n.to_dict() for n in self.notes]
|
|
}
|
|
|
|
@staticmethod
|
|
def from_dict(data):
|
|
ev = Evidence(
|
|
name=data["name"],
|
|
evidence_id=data["evidence_id"],
|
|
description=data.get("description", ""),
|
|
metadata=data.get("metadata", {})
|
|
)
|
|
ev.notes = [Note.from_dict(n) for n in data.get("notes", [])]
|
|
return ev
|
|
|
|
|
|
@dataclass
|
|
class Case:
|
|
case_number: str
|
|
case_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
|
name: str = ""
|
|
investigator: str = ""
|
|
evidence: List[Evidence] = field(default_factory=list)
|
|
notes: List[Note] = field(default_factory=list)
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"case_id": self.case_id,
|
|
"case_number": self.case_number,
|
|
"name": self.name,
|
|
"investigator": self.investigator,
|
|
"evidence": [e.to_dict() for e in self.evidence],
|
|
"notes": [n.to_dict() for n in self.notes]
|
|
}
|
|
|
|
@staticmethod
|
|
def from_dict(data):
|
|
case = Case(
|
|
case_number=data["case_number"],
|
|
case_id=data["case_id"],
|
|
name=data.get("name", ""),
|
|
investigator=data.get("investigator", "")
|
|
)
|
|
case.evidence = [Evidence.from_dict(e) for e in data.get("evidence", [])]
|
|
case.notes = [Note.from_dict(n) for n in data.get("notes", [])]
|
|
return case
|
|
|
|
|
|
__all__ = ['Note', 'Evidence', 'Case', 'TagExtractor', 'IOCExtractor']
|