"""Session management for PentestGPT - persistence and state tracking.""" import json import uuid from dataclasses import dataclass, field from datetime import datetime from enum import Enum from pathlib import Path from typing import Any class SessionStatus(Enum): """Session lifecycle status.""" RUNNING = "running" PAUSED = "paused" COMPLETED = "completed" ERROR = "error" @dataclass class SessionInfo: """Session state - framework agnostic.""" session_id: str target: str created_at: datetime status: SessionStatus = SessionStatus.RUNNING backend_session_id: str | None = None # Backend-specific ID (e.g., Claude session) updated_at: datetime | None = None task: str = "" user_instructions: list[str] = field(default_factory=list) flags_found: list[dict[str, str]] = field(default_factory=list) total_cost_usd: float = 0.0 model: str = "" last_error: str | None = None def to_dict(self) -> dict[str, Any]: """Serialize session to dictionary for JSON storage.""" return { "session_id": self.session_id, "target": self.target, "created_at": self.created_at.isoformat(), "status": self.status.value, "backend_session_id": self.backend_session_id, "updated_at": self.updated_at.isoformat() if self.updated_at else None, "task": self.task, "user_instructions": self.user_instructions, "flags_found": self.flags_found, "total_cost_usd": self.total_cost_usd, "model": self.model, "last_error": self.last_error, } @classmethod def from_dict(cls, data: dict[str, Any]) -> "SessionInfo": """Deserialize session from dictionary.""" return cls( session_id=data["session_id"], target=data["target"], created_at=datetime.fromisoformat(data["created_at"]), status=SessionStatus(data["status"]), backend_session_id=data.get("backend_session_id"), updated_at=( datetime.fromisoformat(data["updated_at"]) if data.get("updated_at") else None ), task=data.get("task", ""), user_instructions=data.get("user_instructions", []), flags_found=data.get("flags_found", []), total_cost_usd=data.get("total_cost_usd", 0.0), model=data.get("model", ""), last_error=data.get("last_error"), ) class SessionStore: """Simple file-based session persistence.""" SESSIONS_DIR = Path.home() / ".pentestgpt" / "sessions" def __init__(self, sessions_dir: Path | None = None): """Initialize session store. Args: sessions_dir: Optional custom sessions directory """ self._sessions_dir = sessions_dir or self.SESSIONS_DIR self._sessions_dir.mkdir(parents=True, exist_ok=True) self._current: SessionInfo | None = None def create(self, target: str, task: str, model: str) -> SessionInfo: """Create a new session. Args: target: Target URL/IP/domain task: Task description model: Model name Returns: New SessionInfo instance """ session = SessionInfo( session_id=str(uuid.uuid4())[:8], target=target, created_at=datetime.now(), task=task, model=model, ) self._current = session self.save() return session @property def current(self) -> SessionInfo | None: """Get current active session.""" return self._current def save(self) -> None: """Save current session to disk.""" if not self._current: return self._current.updated_at = datetime.now() path = self._sessions_dir / f"{self._current.session_id}.json" path.write_text(json.dumps(self._current.to_dict(), indent=2)) def load(self, session_id: str) -> SessionInfo | None: """Load a session by ID. Args: session_id: Session ID to load Returns: SessionInfo if found, None otherwise """ path = self._sessions_dir / f"{session_id}.json" if not path.exists(): return None try: self._current = SessionInfo.from_dict(json.loads(path.read_text())) return self._current except (json.JSONDecodeError, KeyError, ValueError): return None def list_sessions(self, target: str | None = None) -> list[SessionInfo]: """List all sessions, optionally filtered by target. Args: target: Optional target filter Returns: List of SessionInfo, sorted by creation date (newest first) """ sessions = [] for path in self._sessions_dir.glob("*.json"): try: session = SessionInfo.from_dict(json.loads(path.read_text())) if target is None or session.target == target: sessions.append(session) except (json.JSONDecodeError, KeyError, ValueError): continue return sorted(sessions, key=lambda s: s.created_at, reverse=True) def get_latest(self, target: str | None = None) -> SessionInfo | None: """Get the most recent session. Args: target: Optional target filter Returns: Most recent SessionInfo if any exist """ sessions = self.list_sessions(target) return sessions[0] if sessions else None def delete(self, session_id: str) -> bool: """Delete a session by ID. Args: session_id: Session ID to delete Returns: True if deleted, False if not found """ path = self._sessions_dir / f"{session_id}.json" if path.exists(): path.unlink() if self._current and self._current.session_id == session_id: self._current = None return True return False # Convenience methods for updating current session def update_status(self, status: SessionStatus) -> None: """Update current session status.""" if self._current: self._current.status = status self.save() def add_instruction(self, instruction: str) -> None: """Add a user instruction to current session.""" if self._current: self._current.user_instructions.append(instruction) self.save() def add_flag(self, flag: str, context: str) -> None: """Add a found flag to current session.""" if self._current: self._current.flags_found.append({"flag": flag, "context": context}) self.save() def set_backend_session_id(self, backend_id: str) -> None: """Set the backend-specific session ID.""" if self._current: self._current.backend_session_id = backend_id self.save() def add_cost(self, cost: float) -> None: """Add to total cost.""" if self._current: self._current.total_cost_usd += cost self.save() def set_error(self, error: str) -> None: """Set last error.""" if self._current: self._current.last_error = error self.save()