351 lines
11 KiB
Python
351 lines
11 KiB
Python
|
"""
|
||
|
|
AUTARCH Penetration Testing Tree (PTT)
|
||
|
|
Hierarchical task tracker for structured penetration testing workflows.
|
||
|
|
Based on PentestGPT's USENIX paper methodology.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import uuid
|
||
|
|
from enum import Enum
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import Optional, List, Dict, Any
|
||
|
|
|
||
|
|
|
||
|
|
class NodeStatus(Enum):
|
||
|
|
TODO = "todo"
|
||
|
|
IN_PROGRESS = "in_progress"
|
||
|
|
COMPLETED = "completed"
|
||
|
|
NOT_APPLICABLE = "not_applicable"
|
||
|
|
|
||
|
|
|
||
|
|
class PTTNodeType(Enum):
|
||
|
|
RECONNAISSANCE = "reconnaissance"
|
||
|
|
INITIAL_ACCESS = "initial_access"
|
||
|
|
PRIVILEGE_ESCALATION = "privilege_escalation"
|
||
|
|
LATERAL_MOVEMENT = "lateral_movement"
|
||
|
|
PERSISTENCE = "persistence"
|
||
|
|
CREDENTIAL_ACCESS = "credential_access"
|
||
|
|
EXFILTRATION = "exfiltration"
|
||
|
|
CUSTOM = "custom"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class PTTNode:
|
||
|
|
"""A single node in the Penetration Testing Tree."""
|
||
|
|
id: str
|
||
|
|
label: str
|
||
|
|
node_type: PTTNodeType
|
||
|
|
status: NodeStatus = NodeStatus.TODO
|
||
|
|
parent_id: Optional[str] = None
|
||
|
|
children: List[str] = field(default_factory=list)
|
||
|
|
details: str = ""
|
||
|
|
tool_output: Optional[str] = None
|
||
|
|
findings: List[str] = field(default_factory=list)
|
||
|
|
priority: int = 3
|
||
|
|
created_at: str = ""
|
||
|
|
updated_at: str = ""
|
||
|
|
|
||
|
|
def __post_init__(self):
|
||
|
|
now = datetime.now().isoformat()
|
||
|
|
if not self.created_at:
|
||
|
|
self.created_at = now
|
||
|
|
if not self.updated_at:
|
||
|
|
self.updated_at = now
|
||
|
|
|
||
|
|
def to_dict(self) -> dict:
|
||
|
|
return {
|
||
|
|
'id': self.id,
|
||
|
|
'label': self.label,
|
||
|
|
'node_type': self.node_type.value,
|
||
|
|
'status': self.status.value,
|
||
|
|
'parent_id': self.parent_id,
|
||
|
|
'children': self.children.copy(),
|
||
|
|
'details': self.details,
|
||
|
|
'tool_output': self.tool_output,
|
||
|
|
'findings': self.findings.copy(),
|
||
|
|
'priority': self.priority,
|
||
|
|
'created_at': self.created_at,
|
||
|
|
'updated_at': self.updated_at,
|
||
|
|
}
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_dict(cls, data: dict) -> 'PTTNode':
|
||
|
|
return cls(
|
||
|
|
id=data['id'],
|
||
|
|
label=data['label'],
|
||
|
|
node_type=PTTNodeType(data['node_type']),
|
||
|
|
status=NodeStatus(data['status']),
|
||
|
|
parent_id=data.get('parent_id'),
|
||
|
|
children=data.get('children', []),
|
||
|
|
details=data.get('details', ''),
|
||
|
|
tool_output=data.get('tool_output'),
|
||
|
|
findings=data.get('findings', []),
|
||
|
|
priority=data.get('priority', 3),
|
||
|
|
created_at=data.get('created_at', ''),
|
||
|
|
updated_at=data.get('updated_at', ''),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# Status display symbols
|
||
|
|
_STATUS_SYMBOLS = {
|
||
|
|
NodeStatus.TODO: '[ ]',
|
||
|
|
NodeStatus.IN_PROGRESS: '[~]',
|
||
|
|
NodeStatus.COMPLETED: '[x]',
|
||
|
|
NodeStatus.NOT_APPLICABLE: '[-]',
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class PentestTree:
|
||
|
|
"""Penetration Testing Tree - hierarchical task tracker."""
|
||
|
|
|
||
|
|
def __init__(self, target: str):
|
||
|
|
self.target = target
|
||
|
|
self.nodes: Dict[str, PTTNode] = {}
|
||
|
|
self.root_nodes: List[str] = []
|
||
|
|
now = datetime.now().isoformat()
|
||
|
|
self.created_at = now
|
||
|
|
self.updated_at = now
|
||
|
|
|
||
|
|
def add_node(
|
||
|
|
self,
|
||
|
|
label: str,
|
||
|
|
node_type: PTTNodeType,
|
||
|
|
parent_id: Optional[str] = None,
|
||
|
|
details: str = "",
|
||
|
|
priority: int = 3,
|
||
|
|
status: NodeStatus = NodeStatus.TODO,
|
||
|
|
) -> str:
|
||
|
|
"""Add a node to the tree. Returns the new node's ID."""
|
||
|
|
node_id = str(uuid.uuid4())[:8]
|
||
|
|
node = PTTNode(
|
||
|
|
id=node_id,
|
||
|
|
label=label,
|
||
|
|
node_type=node_type,
|
||
|
|
status=status,
|
||
|
|
parent_id=parent_id,
|
||
|
|
details=details,
|
||
|
|
priority=priority,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.nodes[node_id] = node
|
||
|
|
|
||
|
|
if parent_id and parent_id in self.nodes:
|
||
|
|
self.nodes[parent_id].children.append(node_id)
|
||
|
|
elif parent_id is None:
|
||
|
|
self.root_nodes.append(node_id)
|
||
|
|
|
||
|
|
self.updated_at = datetime.now().isoformat()
|
||
|
|
return node_id
|
||
|
|
|
||
|
|
def update_node(
|
||
|
|
self,
|
||
|
|
node_id: str,
|
||
|
|
status: Optional[NodeStatus] = None,
|
||
|
|
details: Optional[str] = None,
|
||
|
|
tool_output: Optional[str] = None,
|
||
|
|
findings: Optional[List[str]] = None,
|
||
|
|
priority: Optional[int] = None,
|
||
|
|
label: Optional[str] = None,
|
||
|
|
) -> bool:
|
||
|
|
"""Update a node's properties. Returns True if found and updated."""
|
||
|
|
node = self.nodes.get(node_id)
|
||
|
|
if not node:
|
||
|
|
return False
|
||
|
|
|
||
|
|
if status is not None:
|
||
|
|
node.status = status
|
||
|
|
if details is not None:
|
||
|
|
node.details = details
|
||
|
|
if tool_output is not None:
|
||
|
|
node.tool_output = tool_output
|
||
|
|
if findings is not None:
|
||
|
|
node.findings.extend(findings)
|
||
|
|
if priority is not None:
|
||
|
|
node.priority = priority
|
||
|
|
if label is not None:
|
||
|
|
node.label = label
|
||
|
|
|
||
|
|
node.updated_at = datetime.now().isoformat()
|
||
|
|
self.updated_at = node.updated_at
|
||
|
|
return True
|
||
|
|
|
||
|
|
def delete_node(self, node_id: str) -> bool:
|
||
|
|
"""Delete a node and all its children recursively."""
|
||
|
|
node = self.nodes.get(node_id)
|
||
|
|
if not node:
|
||
|
|
return False
|
||
|
|
|
||
|
|
# Recursively delete children
|
||
|
|
for child_id in node.children.copy():
|
||
|
|
self.delete_node(child_id)
|
||
|
|
|
||
|
|
# Remove from parent's children list
|
||
|
|
if node.parent_id and node.parent_id in self.nodes:
|
||
|
|
parent = self.nodes[node.parent_id]
|
||
|
|
if node_id in parent.children:
|
||
|
|
parent.children.remove(node_id)
|
||
|
|
|
||
|
|
# Remove from root nodes if applicable
|
||
|
|
if node_id in self.root_nodes:
|
||
|
|
self.root_nodes.remove(node_id)
|
||
|
|
|
||
|
|
del self.nodes[node_id]
|
||
|
|
self.updated_at = datetime.now().isoformat()
|
||
|
|
return True
|
||
|
|
|
||
|
|
def get_node(self, node_id: str) -> Optional[PTTNode]:
|
||
|
|
return self.nodes.get(node_id)
|
||
|
|
|
||
|
|
def get_next_todo(self) -> Optional[PTTNode]:
|
||
|
|
"""Get the highest priority TODO node."""
|
||
|
|
todos = [n for n in self.nodes.values() if n.status == NodeStatus.TODO]
|
||
|
|
if not todos:
|
||
|
|
return None
|
||
|
|
return min(todos, key=lambda n: n.priority)
|
||
|
|
|
||
|
|
def get_all_by_status(self, status: NodeStatus) -> List[PTTNode]:
|
||
|
|
return [n for n in self.nodes.values() if n.status == status]
|
||
|
|
|
||
|
|
def get_subtree(self, node_id: str) -> List[PTTNode]:
|
||
|
|
"""Get all nodes in a subtree (including the root)."""
|
||
|
|
node = self.nodes.get(node_id)
|
||
|
|
if not node:
|
||
|
|
return []
|
||
|
|
result = [node]
|
||
|
|
for child_id in node.children:
|
||
|
|
result.extend(self.get_subtree(child_id))
|
||
|
|
return result
|
||
|
|
|
||
|
|
def find_node_by_label(self, label: str) -> Optional[PTTNode]:
|
||
|
|
"""Find a node by label (case-insensitive partial match)."""
|
||
|
|
label_lower = label.lower()
|
||
|
|
for node in self.nodes.values():
|
||
|
|
if label_lower in node.label.lower():
|
||
|
|
return node
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_stats(self) -> Dict[str, int]:
|
||
|
|
"""Get tree statistics."""
|
||
|
|
stats = {'total': len(self.nodes)}
|
||
|
|
for status in NodeStatus:
|
||
|
|
stats[status.value] = len(self.get_all_by_status(status))
|
||
|
|
return stats
|
||
|
|
|
||
|
|
def render_text(self) -> str:
|
||
|
|
"""Render full tree as indented text for terminal display."""
|
||
|
|
if not self.root_nodes:
|
||
|
|
return " (empty tree)"
|
||
|
|
|
||
|
|
lines = [f"Target: {self.target}"]
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
for root_id in self.root_nodes:
|
||
|
|
self._render_node(root_id, lines, indent=0)
|
||
|
|
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
def _render_node(self, node_id: str, lines: List[str], indent: int):
|
||
|
|
node = self.nodes.get(node_id)
|
||
|
|
if not node:
|
||
|
|
return
|
||
|
|
|
||
|
|
prefix = " " * indent
|
||
|
|
symbol = _STATUS_SYMBOLS.get(node.status, '[ ]')
|
||
|
|
priority_str = f" P{node.priority}" if node.priority != 3 else ""
|
||
|
|
lines.append(f"{prefix}{symbol} {node.label}{priority_str}")
|
||
|
|
|
||
|
|
if node.findings:
|
||
|
|
for finding in node.findings[:3]:
|
||
|
|
lines.append(f"{prefix} -> {finding}")
|
||
|
|
|
||
|
|
for child_id in node.children:
|
||
|
|
self._render_node(child_id, lines, indent + 1)
|
||
|
|
|
||
|
|
def render_summary(self) -> str:
|
||
|
|
"""Render compact summary for LLM context injection.
|
||
|
|
Designed to fit within tight token budgets (4096 ctx).
|
||
|
|
Only shows TODO and IN_PROGRESS nodes with minimal detail.
|
||
|
|
"""
|
||
|
|
stats = self.get_stats()
|
||
|
|
lines = [
|
||
|
|
f"Target: {self.target}",
|
||
|
|
f"Nodes: {stats['total']} total, {stats['todo']} todo, "
|
||
|
|
f"{stats['completed']} done, {stats['in_progress']} active",
|
||
|
|
]
|
||
|
|
|
||
|
|
# Show active and todo nodes only
|
||
|
|
active = self.get_all_by_status(NodeStatus.IN_PROGRESS)
|
||
|
|
todos = sorted(
|
||
|
|
self.get_all_by_status(NodeStatus.TODO),
|
||
|
|
key=lambda n: n.priority
|
||
|
|
)
|
||
|
|
|
||
|
|
if active:
|
||
|
|
lines.append("Active:")
|
||
|
|
for n in active:
|
||
|
|
lines.append(f" [{n.id}] {n.label}")
|
||
|
|
|
||
|
|
if todos:
|
||
|
|
lines.append("Todo:")
|
||
|
|
for n in todos[:5]:
|
||
|
|
lines.append(f" [{n.id}] P{n.priority} {n.label}")
|
||
|
|
if len(todos) > 5:
|
||
|
|
lines.append(f" ... and {len(todos) - 5} more")
|
||
|
|
|
||
|
|
# Show recent findings (last 5)
|
||
|
|
all_findings = []
|
||
|
|
for node in self.nodes.values():
|
||
|
|
if node.findings:
|
||
|
|
for f in node.findings:
|
||
|
|
all_findings.append(f)
|
||
|
|
if all_findings:
|
||
|
|
lines.append("Key findings:")
|
||
|
|
for f in all_findings[-5:]:
|
||
|
|
lines.append(f" - {f}")
|
||
|
|
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
def initialize_standard_branches(self):
|
||
|
|
"""Create standard MITRE ATT&CK-aligned top-level branches."""
|
||
|
|
branches = [
|
||
|
|
("Reconnaissance", PTTNodeType.RECONNAISSANCE, 1,
|
||
|
|
"Information gathering and target enumeration"),
|
||
|
|
("Initial Access", PTTNodeType.INITIAL_ACCESS, 2,
|
||
|
|
"Gaining initial foothold on target"),
|
||
|
|
("Privilege Escalation", PTTNodeType.PRIVILEGE_ESCALATION, 3,
|
||
|
|
"Escalating from initial access to higher privileges"),
|
||
|
|
("Lateral Movement", PTTNodeType.LATERAL_MOVEMENT, 4,
|
||
|
|
"Moving to other systems in the network"),
|
||
|
|
("Credential Access", PTTNodeType.CREDENTIAL_ACCESS, 3,
|
||
|
|
"Obtaining credentials and secrets"),
|
||
|
|
("Persistence", PTTNodeType.PERSISTENCE, 5,
|
||
|
|
"Maintaining access to compromised systems"),
|
||
|
|
]
|
||
|
|
|
||
|
|
for label, ntype, priority, details in branches:
|
||
|
|
self.add_node(
|
||
|
|
label=label,
|
||
|
|
node_type=ntype,
|
||
|
|
priority=priority,
|
||
|
|
details=details,
|
||
|
|
)
|
||
|
|
|
||
|
|
def to_dict(self) -> dict:
|
||
|
|
return {
|
||
|
|
'target': self.target,
|
||
|
|
'created_at': self.created_at,
|
||
|
|
'updated_at': self.updated_at,
|
||
|
|
'root_nodes': self.root_nodes.copy(),
|
||
|
|
'nodes': {nid: n.to_dict() for nid, n in self.nodes.items()},
|
||
|
|
}
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_dict(cls, data: dict) -> 'PentestTree':
|
||
|
|
tree = cls(target=data['target'])
|
||
|
|
tree.created_at = data.get('created_at', '')
|
||
|
|
tree.updated_at = data.get('updated_at', '')
|
||
|
|
tree.root_nodes = data.get('root_nodes', [])
|
||
|
|
for nid, ndata in data.get('nodes', {}).items():
|
||
|
|
tree.nodes[nid] = PTTNode.from_dict(ndata)
|
||
|
|
return tree
|