Skills

Skills — Memory & Knowledge System

PentestAI has a layered memory system inspired by how experienced pentesters work: remembering techniques from previous engagements, building a knowledge graph about the target, and retrieving relevant information when needed.

---

Memory Architecture

┌─────────────────────────────────────────────────┐
│                  MemoryManager                   │
│                 (Singleton)                      │
└────┬────────────┬────────────┬──────────────────┘
     │            │            │
     ▼            ▼            ▼
  Session     Knowledge     Learning
  (current)    Graph        Engine
     │            │            │
     │         SQLite +      Pattern
     │         Neo4j*        Matching
     │
     ├── Vector Store (TF-IDF + Qdrant*)
     ├── RAG Pipeline
     └── CVE Feed (NVD)

* = optional, requires an external service

---

Session Memory

File: memory/session.py

Storage for the state of the engagement currently in progress. All target information is stored here during the session and persisted to disk.

Stored Data

@dataclass
class Session:
    session_id: str           # 8-character UUID
    name: str                 # engagement name
    primary_target: str       # primary IP/hostname
    services: list[ServiceInfo]    # open ports + service info
    findings: list[Finding]        # vulnerabilities found
    endpoints: set[str]            # discovered URL endpoints
    technologies: set[str]         # technology stack
    command_history: list[CommandRecord]
    notes: list[str]
    scope_data: dict               # for restoring scope on load

Finding

@dataclass
class Finding:
    id: str                   # automatic UUID
    title: str                # vulnerability name
    severity: str             # Critical/High/Medium/Low/Info
    target: str               # vulnerable host/URL
    tool: str                 # tool that found it
    description: str          # detailed explanation
    evidence: str             # evidence (tool output, request/response)
    recommendation: str       # remediation approach
    tags: list[str]           # e.g. ["sqli", "owasp-a03"]
    timestamp: str

ServiceInfo

@dataclass
class ServiceInfo:
    host: str
    port: int
    protocol: str      # tcp/udp
    service: str       # ssh/http/ftp/etc
    version: str       # software version
    banner: str        # raw banner
    state: str         # open/filtered
    notes: str

Usage

from memory.session import Session, Finding, ServiceInfo

# Create a session
s = Session(name="htb-forest")
s.set_target("10.10.11.162")

# Add a service
s.add_service(ServiceInfo(
    host="10.10.11.162", port=389, protocol="tcp",
    service="ldap", version="AD LDAP"
))

# Add a finding
s.add_finding(Finding(
    title="LDAP Anonymous Bind",
    severity="High",
    target="ldap://10.10.11.162",
    tool="nmap",
    description="LDAP allows anonymous bind, exposing directory info",
    recommendation="Disable anonymous LDAP access",
    tags=["ldap", "recon"],
))

# Save to disk
path = s.save()

# Load back
s2 = Session.load(path)

# Context for the LLM
context = s.get_llm_context()

---

Knowledge Graph

File: memory/graph/knowledge_graph.py

A graph representing the relationships between hosts, services, endpoints, vulnerabilities, credentials, and technologies within an engagement.

Node Types

Type	Example
HOST	10.10.11.162
SERVICE	10.10.11.162:389/tcp/ldap
ENDPOINT	http://10.10.11.162/api/users
VULNERABILITY	CVE-2020-1472 (ZeroLogon)
CREDENTIAL	svc-alfresco:s3rvice
TECHNOLOGY	Windows Server 2016
ATTACK_PATH	anon_ldap → user_enum → kerberoast

Usage

from memory.graph.knowledge_graph import KnowledgeGraph

g = KnowledgeGraph()  # default: SQLite in sessions/

# Add nodes
host = g.add_host("10.10.11.162")
svc  = g.add_service("10.10.11.162", 389, "tcp", "ldap")
vuln = g.add_vulnerability(
    svc.node_id,
    cve_id="CVE-2020-1472",
    title="ZeroLogon",
    cvss=10.0,
    severity="CRITICAL",
)

# Link attack path
g.link_attack_path(svc.node_id, vuln.node_id, "exploitation")

# Query
paths = g.find_attack_paths("10.10.11.162")
for path in paths:
    print(" → ".join(path))

# Stats
stats = g.stats()
print(f"Nodes: {stats['nodes']}, Edges: {stats['edges']}")

# Import from a session
g.import_from_session(session)

Backend

Defaults to SQLite (included, no setup needed).

For large engagements or complex graph analysis, you can upgrade to Neo4j:

# .env
NEO4J_URI=bolt://localhost:7687
NEO4J_USER=neo4j
NEO4J_PASS=pentestai

# Start Neo4j via Docker
docker run -d \
  -e NEO4J_AUTH=neo4j/pentestai \
  -p 7474:7474 -p 7687:7687 \
  neo4j:5-community

In the Compose stack, Neo4j is already available without additional setup:

docker compose --profile msf up -d neo4j
# from inside the pentestai container, use the service hostname:
# NEO4J_URI=bolt://neo4j:7687

---

Vector Store

File: memory/vector/store.py

Semantic storage and search for findings, techniques, CVEs, and other knowledge.

Usage

from memory.vector.store import VectorStore

store = VectorStore()

# Add a document
store.add_document(
    "SQL injection via UNION SELECT in login parameter allows auth bypass",
    collection="findings",
    metadata={"severity": "High", "session_id": "abc123"},
)

# Search
results = store.search("sql injection authentication bypass", top_k=5)
for r in results:
    print(f"[{r.score:.2f}] {r.document.text[:80]}")

Collections

Collection	Contents
findings	Findings from all sessions
techniques	Successful pentest techniques
cve	CVE entries from NVD
payloads	Effective payloads

Backend

Defaults to in-memory TF-IDF (included).

For larger scale, upgrade to Qdrant:

QDRANT_URL=http://localhost:6333

docker run -d -p 6333:6333 qdrant/qdrant:latest

In the Compose stack: docker compose --profile msf up -d qdrant, then set QDRANT_URL=http://qdrant:6333 from inside the container.

---

CVE Feed

File: memory/knowledge/cve_feed.py

Integration with the National Vulnerability Database (NVD) to enrich findings with CVE data.

Usage

from memory.knowledge.cve_feed import CVEFeed

feed = CVEFeed()

# Search CVEs by keyword
cves = feed.search("log4j jndi injection", top_k=5)
for cve in cves:
    print(f"{cve.cve_id} [{cve.severity}] CVSS:{cve.cvss_v3}")
    print(f"  {cve.description[:100]}")

# Get a specific CVE
cve = feed.get_cve("CVE-2021-44228")
print(cve.description)

# Get CVEs relevant to the technologies found
cves = feed.get_relevant_for_technologies(
    ["Apache 2.4.49", "Log4j 2.14.1", "Spring Framework 5.3"],
)

# Refresh the cache from NVD
count = feed.refresh_recent(days=7)
print(f"Indexed {count} CVEs")

Rate Limiting

Without an NVD API key: 10 requests/minute
With an NVD API key: 50 requests/minute

Register for a free key at nvd.nist.gov/developers:

NVD_API_KEY=your_nvd_api_key

---

Learning Engine

File: memory/learning.py

⚠️ Disabled by design. get_recommendations() always returns [] because heuristic cross-session matching created false confidence. The module and its storage remain for a future fine-tuned replacement. See STATUS.md.

Engine for learning from previous engagements and recommending techniques based on the target profile.

How It Works

1. After an engagement ends, learn_from_session() extracts patterns
2. Patterns stored: technology stack, services, successful tools, findings
3. For a new engagement, get_recommendations() matches based on similarity score
4. Recommendations = tools and techniques that succeeded on similar targets

Usage

from memory.learning import LearningEngine

engine = LearningEngine()

# After the engagement ends
engine.learn_from_session(session)

# For a new engagement
recommendations = engine.get_recommendations(
    services=["http", "ssh", "smb"],
    technologies=["Windows Server 2019", "IIS"],
)
for rec in recommendations[:5]:
    print(f"  Sequence: {' → '.join(rec['sequence'])}")
    print(f"  Similarity: {rec['similarity']:.0%}")
    print(f"  Past findings: {', '.join(rec['findings'][:3])}")

CLI

memory learn                  # learn from the current session
memory recommend              # show recommendations for the target

---

RAG Pipeline

File: memory/rag/pipeline.py

A Retrieval-Augmented Generation pipeline that combines vector search and the CVE feed to provide relevant context to the LLM.

Usage

from memory.rag.pipeline import RAGPipeline

rag = RAGPipeline(vector_store=store, cve_feed=feed)

# Retrieve relevant context
context = rag.retrieve(
    "apache struts remote code execution",
    top_k=5,
)

# Use in the LLM prompt
if context.has_content():
    prompt = f"""
{context.to_context_string(max_chars=2000)}

Given the above context about vulnerabilities, analyze...
"""

---

MemoryManager

File: memory/manager.py

A singleton that coordinates all memory components.

Usage

from memory.manager import MemoryManager

# Singleton — always the same instance
mgr = MemoryManager.get()

# Import the session into all stores
mgr.import_session(session)

# Search the knowledge base
results = mgr.search_knowledge("kerberoasting techniques", top_k=10)
print(results["findings"])    # past findings
print(results["techniques"])  # relevant techniques
print(results["cves"])        # relevant CVEs

# Recommendations
recs = mgr.recommend_next_steps(session)

# Context for the LLM
enhanced_prompt = mgr.enrich_prompt(base_prompt, session)

# Status of all systems
status = mgr.status()
print(f"Graph: {status['graph']['nodes']} nodes")
print(f"Vector: {status['vector']['documents']} docs")
print(f"CVE: {status['cve']['total']} entries")

CLI

memory status                 # show the status of all systems
memory import                 # import the current session into the knowledge graph
memory graph 10.10.11.162     # show the attack graph
memory recommend              # technique recommendations
search log4j rce              # search the knowledge base

---

Phase 6 Skills

Multi-Agent Context Sharing

When orchestrate runs, all sub-agents share a single Session object. Findings from ReconSubAgent are immediately available to WebSubAgent and InfraSubAgent without IPC overhead.

PTES Phase Knowledge

_infer_ptes_phase() in reports/ptes.py maps tool names to PTES phases deterministically — no LLM call, no latency.

Plugin Isolation

Plugins installed from the marketplace have their checksum verified and are parsed with ast.parse() before execution — preventing malicious code injection.