mirror of
https://github.com/GH05TCREW/pentestagent.git
synced 2026-03-07 22:33:38 +00:00
feat(workspaces): add unified /workspace lifecycle, target persistence, and workspace-scoped RAG
- Introduce command for CLI and TUI with create/activate, list, info, note, clear, export, import, and help actions - Persist workspace state via marker and enriched (targets, operator notes, last_active_at, last_target) - Restore on workspace activation and sync it to UI banner, agent state, and CLI output - Enforce target normalization and ensure always exists in workspace targets - Route loot output to when a workspace is active - Prefer workspace-local knowledge paths for indexing and RAG resolution - Persist RAG indexes per workspace and load existing indexes before re-indexing - Add deterministic workspace export/import utilities (excluding caches) - Integrate workspace handling into TUI slash commands with modal help screen
This commit is contained in:
50
tests/test_rag_workspace_integration.py
Normal file
50
tests/test_rag_workspace_integration.py
Normal file
@@ -0,0 +1,50 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from pentestagent.workspaces.manager import WorkspaceManager
|
||||
from pentestagent.knowledge.rag import RAGEngine
|
||||
from pentestagent.knowledge.indexer import KnowledgeIndexer
|
||||
|
||||
|
||||
def test_rag_and_indexer_use_workspace(tmp_path, monkeypatch):
|
||||
# Use tmp_path as the project root
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
name = "ws_test"
|
||||
wm.create(name)
|
||||
wm.set_active(name)
|
||||
|
||||
# Create a sample source file in the workspace sources
|
||||
src_dir = tmp_path / "workspaces" / name / "knowledge" / "sources"
|
||||
src_dir.mkdir(parents=True, exist_ok=True)
|
||||
sample = src_dir / "sample.md"
|
||||
sample.write_text("# Sample\n\nThis is a test knowledge document for RAG indexing.")
|
||||
|
||||
# Ensure KnowledgeIndexer picks up the workspace source when indexing default 'knowledge'
|
||||
ki = KnowledgeIndexer()
|
||||
docs, result = ki.index_directory(Path("knowledge"))
|
||||
|
||||
assert result.indexed_files >= 1
|
||||
assert len(docs) >= 1
|
||||
# Ensure the document source path points at the workspace file
|
||||
assert any("workspaces" in d.source and "sample.md" in d.source for d in docs)
|
||||
|
||||
# Now run RAGEngine to build embeddings and verify saved index file appears
|
||||
rag = RAGEngine(use_local_embeddings=True)
|
||||
rag.index()
|
||||
|
||||
emb_path = tmp_path / "workspaces" / name / "knowledge" / "embeddings" / "index.pkl"
|
||||
assert emb_path.exists(), f"Expected saved index at {emb_path}"
|
||||
|
||||
# Ensure RAG engine has documents/chunks loaded
|
||||
assert rag.get_chunk_count() >= 1
|
||||
assert rag.get_document_count() >= 1
|
||||
|
||||
# Now create a new RAGEngine and ensure it loads persisted index automatically
|
||||
rag2 = RAGEngine(use_local_embeddings=True)
|
||||
# If load-on-init doesn't run, calling index() should load from saved file
|
||||
rag2.index()
|
||||
assert rag2.get_chunk_count() >= 1
|
||||
96
tests/test_workspace.py
Normal file
96
tests/test_workspace.py
Normal file
@@ -0,0 +1,96 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from pentestagent.workspaces.manager import WorkspaceManager, WorkspaceError
|
||||
|
||||
|
||||
def test_invalid_workspace_names(tmp_path: Path):
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
bad_names = ["../escape", "name/with/slash", "..", ""]
|
||||
# overlong name
|
||||
bad_names.append("a" * 65)
|
||||
for n in bad_names:
|
||||
with pytest.raises(WorkspaceError):
|
||||
wm.create(n)
|
||||
|
||||
|
||||
def test_create_and_idempotent(tmp_path: Path):
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
name = "eng1"
|
||||
meta = wm.create(name)
|
||||
assert (tmp_path / "workspaces" / name).exists()
|
||||
assert (tmp_path / "workspaces" / name / "meta.yaml").exists()
|
||||
# create again should not raise and should return meta
|
||||
meta2 = wm.create(name)
|
||||
assert meta2["name"] == name
|
||||
|
||||
|
||||
def test_set_get_active(tmp_path: Path):
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
name = "activews"
|
||||
wm.create(name)
|
||||
wm.set_active(name)
|
||||
assert wm.get_active() == name
|
||||
marker = tmp_path / "workspaces" / ".active"
|
||||
assert marker.exists()
|
||||
assert marker.read_text(encoding="utf-8").strip() == name
|
||||
|
||||
|
||||
def test_add_list_remove_targets(tmp_path: Path):
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
name = "targets"
|
||||
wm.create(name)
|
||||
added = wm.add_targets(name, ["192.168.1.1", "192.168.0.0/16", "Example.COM"]) # hostname mixed case
|
||||
# normalized entries
|
||||
assert "192.168.1.1" in added
|
||||
assert "192.168.0.0/16" in added
|
||||
assert "example.com" in added
|
||||
# dedupe
|
||||
added2 = wm.add_targets(name, ["192.168.1.1", "example.com"])
|
||||
assert len(added2) == len(added)
|
||||
# remove
|
||||
after = wm.remove_target(name, "192.168.1.1")
|
||||
assert "192.168.1.1" not in after
|
||||
|
||||
|
||||
def test_persistence_across_instances(tmp_path: Path):
|
||||
wm1 = WorkspaceManager(root=tmp_path)
|
||||
name = "persist"
|
||||
wm1.create(name)
|
||||
wm1.add_targets(name, ["10.0.0.1", "host.local"])
|
||||
|
||||
# new manager instance reads from disk
|
||||
wm2 = WorkspaceManager(root=tmp_path)
|
||||
targets = wm2.list_targets(name)
|
||||
assert "10.0.0.1" in targets
|
||||
assert "host.local" in targets
|
||||
|
||||
|
||||
def test_last_target_persistence(tmp_path: Path):
|
||||
wm = WorkspaceManager(root=tmp_path)
|
||||
a = "wsA"
|
||||
b = "wsB"
|
||||
wm.create(a)
|
||||
wm.create(b)
|
||||
|
||||
t1 = "192.168.0.4"
|
||||
t2 = "192.168.0.165"
|
||||
|
||||
# set last target on workspace A and B
|
||||
norm1 = wm.set_last_target(a, t1)
|
||||
norm2 = wm.set_last_target(b, t2)
|
||||
|
||||
# persisted in meta
|
||||
assert wm.get_meta_field(a, "last_target") == norm1
|
||||
assert wm.get_meta_field(b, "last_target") == norm2
|
||||
|
||||
# targets list contains the last target
|
||||
assert norm1 in wm.list_targets(a)
|
||||
assert norm2 in wm.list_targets(b)
|
||||
|
||||
# new manager instance still sees last_target
|
||||
wm2 = WorkspaceManager(root=tmp_path)
|
||||
assert wm2.get_meta_field(a, "last_target") == norm1
|
||||
assert wm2.get_meta_field(b, "last_target") == norm2
|
||||
Reference in New Issue
Block a user