Autarch Will Control The Internet
This commit is contained in:
DigiJ
279
core/pentest_session.py
Normal file
279
core/pentest_session.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""
|
||||
AUTARCH Pentest Session Manager
|
||||
Save and resume penetration testing sessions with full state persistence.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
|
||||
from .pentest_tree import PentestTree, NodeStatus
|
||||
|
||||
|
||||
class PentestSessionState(Enum):
|
||||
IDLE = "idle"
|
||||
RUNNING = "running"
|
||||
PAUSED = "paused"
|
||||
COMPLETED = "completed"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionEvent:
|
||||
"""A single event in the session timeline."""
|
||||
timestamp: str
|
||||
event_type: str
|
||||
data: dict
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
'timestamp': self.timestamp,
|
||||
'event_type': self.event_type,
|
||||
'data': self.data,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> 'SessionEvent':
|
||||
return cls(
|
||||
timestamp=data['timestamp'],
|
||||
event_type=data['event_type'],
|
||||
data=data.get('data', {}),
|
||||
)
|
||||
|
||||
|
||||
class PentestSession:
|
||||
"""Manages a single penetration testing session."""
|
||||
|
||||
@classmethod
|
||||
def _get_dir(cls):
|
||||
from core.paths import get_data_dir
|
||||
d = get_data_dir() / "pentest_sessions"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d
|
||||
|
||||
def __init__(self, target: str, session_id: str = None):
|
||||
self.session_id = session_id or self._generate_id(target)
|
||||
self.target = target
|
||||
self.state = PentestSessionState.IDLE
|
||||
self.tree = PentestTree(target)
|
||||
self.events: List[SessionEvent] = []
|
||||
self.findings: List[Dict[str, Any]] = []
|
||||
self.pipeline_history: List[dict] = []
|
||||
self.notes: str = ""
|
||||
self.step_count: int = 0
|
||||
now = datetime.now().isoformat()
|
||||
self.created_at = now
|
||||
self.updated_at = now
|
||||
|
||||
@staticmethod
|
||||
def _generate_id(target: str) -> str:
|
||||
"""Generate a session ID from target and timestamp."""
|
||||
safe = re.sub(r'[^a-zA-Z0-9]', '_', target)[:30]
|
||||
ts = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
return f"{safe}_{ts}"
|
||||
|
||||
def start(self):
|
||||
"""Initialize a new session."""
|
||||
self.state = PentestSessionState.RUNNING
|
||||
self.tree.initialize_standard_branches()
|
||||
self.log_event('state_change', {'from': 'idle', 'to': 'running'})
|
||||
self.save()
|
||||
|
||||
def pause(self):
|
||||
"""Pause the session and save state."""
|
||||
prev = self.state.value
|
||||
self.state = PentestSessionState.PAUSED
|
||||
self.log_event('state_change', {'from': prev, 'to': 'paused'})
|
||||
self.save()
|
||||
|
||||
def resume(self):
|
||||
"""Resume a paused session."""
|
||||
prev = self.state.value
|
||||
self.state = PentestSessionState.RUNNING
|
||||
self.log_event('state_change', {'from': prev, 'to': 'running'})
|
||||
self.save()
|
||||
|
||||
def complete(self, summary: str = ""):
|
||||
"""Mark session as completed."""
|
||||
prev = self.state.value
|
||||
self.state = PentestSessionState.COMPLETED
|
||||
self.log_event('state_change', {
|
||||
'from': prev,
|
||||
'to': 'completed',
|
||||
'summary': summary,
|
||||
})
|
||||
self.save()
|
||||
|
||||
def set_error(self, error_msg: str):
|
||||
"""Mark session as errored."""
|
||||
prev = self.state.value
|
||||
self.state = PentestSessionState.ERROR
|
||||
self.log_event('state_change', {
|
||||
'from': prev,
|
||||
'to': 'error',
|
||||
'error': error_msg,
|
||||
})
|
||||
self.save()
|
||||
|
||||
def log_event(self, event_type: str, data: dict):
|
||||
"""Log an event to the session timeline."""
|
||||
event = SessionEvent(
|
||||
timestamp=datetime.now().isoformat(),
|
||||
event_type=event_type,
|
||||
data=data,
|
||||
)
|
||||
self.events.append(event)
|
||||
self.updated_at = event.timestamp
|
||||
|
||||
def log_pipeline_result(self, parsed: str, reasoning: str, actions: list):
|
||||
"""Log a pipeline execution cycle."""
|
||||
self.pipeline_history.append({
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'step': self.step_count,
|
||||
'parsed_input': parsed,
|
||||
'reasoning': reasoning,
|
||||
'generated_actions': actions,
|
||||
})
|
||||
self.step_count += 1
|
||||
|
||||
def add_finding(self, title: str, description: str,
|
||||
severity: str = "medium", node_id: str = None):
|
||||
"""Add a key finding."""
|
||||
self.findings.append({
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'severity': severity,
|
||||
'title': title,
|
||||
'description': description,
|
||||
'node_id': node_id,
|
||||
})
|
||||
|
||||
def save(self) -> str:
|
||||
"""Save session to JSON file. Returns filepath."""
|
||||
self._get_dir().mkdir(parents=True, exist_ok=True)
|
||||
filepath = self._get_dir() / f"{self.session_id}.json"
|
||||
|
||||
data = {
|
||||
'session_id': self.session_id,
|
||||
'target': self.target,
|
||||
'state': self.state.value,
|
||||
'created_at': self.created_at,
|
||||
'updated_at': self.updated_at,
|
||||
'notes': self.notes,
|
||||
'step_count': self.step_count,
|
||||
'tree': self.tree.to_dict(),
|
||||
'events': [e.to_dict() for e in self.events],
|
||||
'findings': self.findings,
|
||||
'pipeline_history': self.pipeline_history,
|
||||
}
|
||||
|
||||
with open(filepath, 'w') as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
return str(filepath)
|
||||
|
||||
@classmethod
|
||||
def load_session(cls, session_id: str) -> 'PentestSession':
|
||||
"""Load a session from file."""
|
||||
filepath = cls._get_dir() / f"{session_id}.json"
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"Session not found: {session_id}")
|
||||
|
||||
with open(filepath, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
session = cls(target=data['target'], session_id=data['session_id'])
|
||||
session.state = PentestSessionState(data['state'])
|
||||
session.created_at = data['created_at']
|
||||
session.updated_at = data['updated_at']
|
||||
session.notes = data.get('notes', '')
|
||||
session.step_count = data.get('step_count', 0)
|
||||
session.tree = PentestTree.from_dict(data['tree'])
|
||||
session.events = [SessionEvent.from_dict(e) for e in data.get('events', [])]
|
||||
session.findings = data.get('findings', [])
|
||||
session.pipeline_history = data.get('pipeline_history', [])
|
||||
return session
|
||||
|
||||
@classmethod
|
||||
def list_sessions(cls) -> List[Dict[str, Any]]:
|
||||
"""List all saved sessions with summary info."""
|
||||
cls._get_dir().mkdir(parents=True, exist_ok=True)
|
||||
sessions = []
|
||||
for f in sorted(cls._get_dir().glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True):
|
||||
try:
|
||||
with open(f, 'r') as fh:
|
||||
data = json.load(fh)
|
||||
stats = {}
|
||||
if 'tree' in data and 'nodes' in data['tree']:
|
||||
nodes = data['tree']['nodes']
|
||||
stats = {
|
||||
'total': len(nodes),
|
||||
'todo': sum(1 for n in nodes.values() if n.get('status') == 'todo'),
|
||||
'completed': sum(1 for n in nodes.values() if n.get('status') == 'completed'),
|
||||
}
|
||||
sessions.append({
|
||||
'session_id': data['session_id'],
|
||||
'target': data['target'],
|
||||
'state': data['state'],
|
||||
'created': data['created_at'],
|
||||
'updated': data['updated_at'],
|
||||
'steps': data.get('step_count', 0),
|
||||
'findings': len(data.get('findings', [])),
|
||||
'tree_stats': stats,
|
||||
})
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
continue
|
||||
return sessions
|
||||
|
||||
def delete(self) -> bool:
|
||||
"""Delete this session's file."""
|
||||
filepath = self._get_dir() / f"{self.session_id}.json"
|
||||
if filepath.exists():
|
||||
filepath.unlink()
|
||||
return True
|
||||
return False
|
||||
|
||||
def export_report(self) -> str:
|
||||
"""Generate a text summary report of the session."""
|
||||
stats = self.tree.get_stats()
|
||||
lines = [
|
||||
"=" * 60,
|
||||
"AUTARCH Pentest Session Report",
|
||||
"=" * 60,
|
||||
f"Target: {self.target}",
|
||||
f"Session: {self.session_id}",
|
||||
f"State: {self.state.value}",
|
||||
f"Started: {self.created_at}",
|
||||
f"Updated: {self.updated_at}",
|
||||
f"Steps: {self.step_count}",
|
||||
"",
|
||||
"--- Task Tree ---",
|
||||
f"Total nodes: {stats['total']}",
|
||||
f" Completed: {stats.get('completed', 0)}",
|
||||
f" Todo: {stats.get('todo', 0)}",
|
||||
f" Active: {stats.get('in_progress', 0)}",
|
||||
f" N/A: {stats.get('not_applicable', 0)}",
|
||||
"",
|
||||
self.tree.render_text(),
|
||||
"",
|
||||
]
|
||||
|
||||
if self.findings:
|
||||
lines.append("--- Findings ---")
|
||||
for i, f in enumerate(self.findings, 1):
|
||||
sev = f.get('severity', 'medium').upper()
|
||||
lines.append(f" [{i}] [{sev}] {f['title']}")
|
||||
lines.append(f" {f['description']}")
|
||||
lines.append("")
|
||||
|
||||
if self.notes:
|
||||
lines.append("--- Notes ---")
|
||||
lines.append(self.notes)
|
||||
lines.append("")
|
||||
|
||||
lines.append("=" * 60)
|
||||
return "\n".join(lines)
|
||||
Reference in New Issue
Block a user