From 672fbb3550a69216a493d1083098c1fa8a2586e0 Mon Sep 17 00:00:00 2001 From: harmj0y Date: Thu, 6 Nov 2025 18:30:11 -0800 Subject: [PATCH 01/14] Initial Chatbot commit - Initial Chatbot commit --- compose.yaml | 3 + env.example | 13 + infra/postgres/01-schema.sql | 34 +- projects/agents/Dockerfile | 17 +- projects/agents/agents/main.py | 17 + projects/agents/agents/mcp/tools.yaml | 407 ++++++++++++++++++ projects/agents/agents/tasks/chatbot.py | 313 ++++++++++++++ projects/frontend/src/App.jsx | 11 +- .../src/components/Chatbot/ChatbotPage.jsx | 291 +++++++++++++ .../src/components/Chatbot/ExampleQueries.jsx | 53 +++ .../src/components/Chatbot/MessageBubble.jsx | 91 ++++ .../src/components/Chatbot/QueryModal.jsx | 117 +++++ projects/web_api/poetry.lock | 2 +- projects/web_api/pyproject.toml | 1 + projects/web_api/web_api/main.py | 57 ++- projects/web_api/web_api/models/requests.py | 16 + 16 files changed, 1437 insertions(+), 6 deletions(-) create mode 100644 projects/agents/agents/mcp/tools.yaml create mode 100644 projects/agents/agents/tasks/chatbot.py create mode 100644 projects/frontend/src/components/Chatbot/ChatbotPage.jsx create mode 100644 projects/frontend/src/components/Chatbot/ExampleQueries.jsx create mode 100644 projects/frontend/src/components/Chatbot/MessageBubble.jsx create mode 100644 projects/frontend/src/components/Chatbot/QueryModal.jsx diff --git a/compose.yaml b/compose.yaml index d0c7e87..2cae72f 100644 --- a/compose.yaml +++ b/compose.yaml @@ -1040,6 +1040,9 @@ services: - NEMESIS_MONITORING=${NEMESIS_MONITORING:-disabled} - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://otel-collector:4317 - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_INSECURE=true + # Chatbot configuration + - CHATBOT_DB_PASSWORD=${CHATBOT_DB_PASSWORD:-chatbot_pass_change_me} + - MCP_MAX_POSTGRES_ROWS=${MCP_MAX_POSTGRES_ROWS:-1000} logging: *logging-config depends_on: postgres: { condition: service_healthy } diff --git a/env.example b/env.example index 1928b65..3d58d4b 100644 --- a/env.example +++ b/env.example @@ -65,3 +65,16 @@ NEMESIS_URL="https://localhost:7443/" # Phoenix UI will be available at http://localhost:6006 # Example: # PHOENIX_ENABLED=true + + +# (Optional) Chatbot database configuration. +# Password for the read-only database user used by the chatbot. +# Example: +# CHATBOT_DB_PASSWORD="chatbot_secure_password" +CHATBOT_DB_PASSWORD="chatbot_pass_change_me" + +# (Optional) Maximum number of rows returned per chatbot query. +# Helps prevent expensive queries and context window issues. +# Example: +# MCP_MAX_POSTGRES_ROWS=500 +MCP_MAX_POSTGRES_ROWS=1000 diff --git a/infra/postgres/01-schema.sql b/infra/postgres/01-schema.sql index 8cd218f..a0f783a 100644 --- a/infra/postgres/01-schema.sql +++ b/infra/postgres/01-schema.sql @@ -837,4 +837,36 @@ CREATE OR REPLACE TRIGGER update_dpapi_domain_backup_keys_updated_at CREATE OR REPLACE TRIGGER update_dpapi_system_credentials_updated_at BEFORE UPDATE ON dpapi.system_credentials FOR EACH ROW - EXECUTE FUNCTION update_updated_at_column(); \ No newline at end of file + EXECUTE FUNCTION update_updated_at_column(); + + +----------------------- +-- CHATBOT READ-ONLY USER +----------------------- +-- Create read-only user for chatbot queries with restricted table access +DO $$ +BEGIN + IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'chatbot_readonly') THEN + CREATE USER chatbot_readonly WITH PASSWORD 'chatbot_pass_change_me'; + END IF; +END +$$; + +-- Grant connection and schema usage +GRANT CONNECT ON DATABASE enrichment TO chatbot_readonly; +GRANT USAGE ON SCHEMA public TO chatbot_readonly; +GRANT USAGE ON SCHEMA chromium TO chatbot_readonly; + +-- Grant SELECT on specific tables only (chatbot-accessible tables) +GRANT SELECT ON files_enriched TO chatbot_readonly; +GRANT SELECT ON enrichments TO chatbot_readonly; +GRANT SELECT ON findings TO chatbot_readonly; +GRANT SELECT ON file_linkings TO chatbot_readonly; +GRANT SELECT ON chromium.cookies TO chatbot_readonly; +GRANT SELECT ON chromium.logins TO chatbot_readonly; + +-- Explicitly revoke write permissions to ensure read-only access +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA public FROM chatbot_readonly; +REVOKE INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA chromium FROM chatbot_readonly; +REVOKE CREATE ON SCHEMA public FROM chatbot_readonly; +REVOKE CREATE ON SCHEMA chromium FROM chatbot_readonly; \ No newline at end of file diff --git a/projects/agents/Dockerfile b/projects/agents/Dockerfile index 8ef5efc..e08b04f 100644 --- a/projects/agents/Dockerfile +++ b/projects/agents/Dockerfile @@ -7,7 +7,17 @@ RUN apt-get update && \ apt-get install -y libpq5 \ gcc libc6-dev curl wget libicu-dev && \ apt-get clean && \ - rm -rf /var/lib/apt/lists/* + rm -rf /var/lib/apt/lists/* && \ + ARCH=$(dpkg --print-architecture) && \ + wget https://go.dev/dl/go1.25.4.linux-${ARCH}.tar.gz && \ + tar -C /usr/local -xzf go1.25.4.linux-${ARCH}.tar.gz && \ + rm go1.25.4.linux-${ARCH}.tar.gz + + +# Install genai-toolbox for chatbot MCP functionality +ENV GOPATH=/go \ + PATH=/usr/local/go/bin:/go/bin:$PATH +RUN go install github.com/googleapis/genai-toolbox@v0.18.0 # Install .NET Runtime @@ -82,6 +92,11 @@ ENV DOTNET_ROOT=/usr/local/dotnet \ PATH=/usr/local/dotnet:$PATH COPY --from=base /usr/local/dotnet /usr/local/dotnet +# Copy genai-toolbox binary from base stage +ENV GOPATH=/go \ + PATH=/go/bin:$PATH +COPY --from=base /go/bin/genai-toolbox /go/bin/genai-toolbox + COPY --from=bundle /venv /venv diff --git a/projects/agents/agents/main.py b/projects/agents/agents/main.py index 7a53b3a..01accec 100644 --- a/projects/agents/agents/main.py +++ b/projects/agents/agents/main.py @@ -677,6 +677,23 @@ def run_report_generator(request: dict): return {"success": False, "error": str(e)} +@app.post("/agents/chatbot/stream") +async def chatbot_stream_endpoint(request: dict): + """Stream chatbot responses for interactive querying.""" + try: + from agents.tasks.chatbot import ChatbotRequest, chatbot_stream + + # Parse and validate request + chatbot_request = ChatbotRequest(**request) + + # Stream the response + return await chatbot_stream(chatbot_request) + + except Exception as e: + logger.exception(message="Error in chatbot streaming") + return {"success": False, "error": str(e)} + + @app.api_route("/healthz", methods=["GET", "HEAD"]) async def health_check(): """Health check endpoint.""" diff --git a/projects/agents/agents/mcp/tools.yaml b/projects/agents/agents/mcp/tools.yaml new file mode 100644 index 0000000..9e0670c --- /dev/null +++ b/projects/agents/agents/mcp/tools.yaml @@ -0,0 +1,407 @@ +# genai-toolbox configuration for Nemesis Chatbot +# Database source definition +sources: + chatbot-db: + kind: postgres + host: ${POSTGRES_HOST:postgres} + port: ${POSTGRES_PORT:5432} + database: ${POSTGRES_DB:enrichment} + user: chatbot_readonly + password: ${CHATBOT_DB_PASSWORD} + +# Custom SQL tools for querying Nemesis data +tools: + # FILES_ENRICHED queries + count-files: + kind: postgres-sql + source: chatbot-db + description: Count total files, optionally filtered by project, agent, source, or extension + parameters: + - name: project + type: string + description: Filter by project name (optional) + - name: agent_id + type: string + description: Filter by agent ID (optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: extension + type: string + description: Filter by file extension (optional) + statement: | + SELECT COUNT(*) as file_count + FROM files_enriched + WHERE ($1::text IS NULL OR project = $1) + AND ($2::text IS NULL OR agent_id = $2) + AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%')) + AND ($4::text IS NULL OR extension = $4) + + search-files: + kind: postgres-sql + source: chatbot-db + description: Search for files by name, path, or extension with optional filters + parameters: + - name: filename_pattern + type: string + description: Search pattern for filename (case-insensitive, optional) + - name: path_pattern + type: string + description: Search pattern for file path (case-insensitive, optional) + - name: extension + type: string + description: Filter by file extension (optional) + - name: project + type: string + description: Filter by project name (optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT object_id, file_name, path, extension, size, magic_type, mime_type, + source, agent_id, project, timestamp + FROM files_enriched + WHERE ($1::text IS NULL OR LOWER(file_name) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(path) LIKE LOWER('%' || $2 || '%')) + AND ($3::text IS NULL OR extension = $3) + AND ($4::text IS NULL OR project = $4) + AND ($5::text IS NULL OR LOWER(source) LIKE LOWER('%' || $5 || '%')) + ORDER BY timestamp DESC + LIMIT LEAST(COALESCE($6, 100), 1000) + + get-file-details: + kind: postgres-sql + source: chatbot-db + description: Get detailed information about a specific file by object_id + parameters: + - name: object_id + type: string + description: The UUID of the file to retrieve + statement: | + SELECT * + FROM files_enriched + WHERE object_id = $1::uuid + + # ENRICHMENTS queries + list-enrichment-modules: + kind: postgres-sql + source: chatbot-db + description: List all unique enrichment module names with their usage counts + statement: | + SELECT module_name, COUNT(*) as usage_count + FROM enrichments + GROUP BY module_name + ORDER BY usage_count DESC + + get-file-enrichments: + kind: postgres-sql + source: chatbot-db + description: Get all enrichment results for a specific file + parameters: + - name: object_id + type: string + description: The UUID of the file + statement: | + SELECT e.enrichment_id, e.module_name, e.result_data, e.created_at, + f.file_name, f.path + FROM enrichments e + JOIN files_enriched f ON e.object_id = f.object_id + WHERE e.object_id = $1::uuid + ORDER BY e.created_at DESC + + search-enrichments-by-module: + kind: postgres-sql + source: chatbot-db + description: Search enrichment results from a specific module + parameters: + - name: module_name + type: string + description: Name of the enrichment module + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT e.object_id, e.module_name, e.result_data, e.created_at, + f.file_name, f.path, f.source + FROM enrichments e + JOIN files_enriched f ON e.object_id = f.object_id + WHERE e.module_name = $1 + ORDER BY e.created_at DESC + LIMIT LEAST(COALESCE($2, 100), 1000) + + # FINDINGS queries + count-findings: + kind: postgres-sql + source: chatbot-db + description: Count findings, optionally filtered by severity, category, or source + parameters: + - name: min_severity + type: integer + description: Minimum severity level (0-10, optional) + - name: category + type: string + description: Filter by finding category (optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + statement: | + SELECT COUNT(*) as finding_count + FROM findings f + JOIN files_enriched fe ON f.object_id = fe.object_id + WHERE ($1::integer IS NULL OR f.severity >= $1) + AND ($2::text IS NULL OR f.category = $2) + AND ($3::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $3 || '%')) + + search-findings: + kind: postgres-sql + source: chatbot-db + description: Search findings with filters for severity, category, and source + parameters: + - name: min_severity + type: integer + description: Minimum severity level (0-10, optional) + - name: category + type: string + description: Filter by finding category (optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT f.finding_id, f.finding_name, f.category, f.severity, f.origin_type, + f.origin_name, f.data, f.created_at, + fe.file_name, fe.path, fe.source, fe.agent_id, fe.project + FROM findings f + JOIN files_enriched fe ON f.object_id = fe.object_id + WHERE ($1::integer IS NULL OR f.severity >= $1) + AND ($2::text IS NULL OR f.category = $2) + AND ($3::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $3 || '%')) + ORDER BY f.severity DESC, f.created_at DESC + LIMIT LEAST(COALESCE($4, 100), 1000) + + get-findings-by-category: + kind: postgres-sql + source: chatbot-db + description: Get aggregated count of findings grouped by category + parameters: + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + statement: | + SELECT f.category, COUNT(*) as count, + AVG(f.severity) as avg_severity, + MAX(f.severity) as max_severity + FROM findings f + JOIN files_enriched fe ON f.object_id = fe.object_id + WHERE ($1::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $1 || '%')) + GROUP BY f.category + ORDER BY count DESC + + get-triaged-findings: + kind: postgres-sql + source: chatbot-db + description: Get findings with triage status (true positive, false positive, etc) + parameters: + - name: triage_value + type: string + description: Filter by triage value (e.g., true_positive, false_positive, optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT f.finding_id, f.finding_name, f.category, f.severity, + fth.value as triage_value, fth.explanation, fth.timestamp as triage_timestamp, + fe.file_name, fe.path, fe.source + FROM findings f + JOIN findings_triage_history fth ON f.triage_id = fth.id + JOIN files_enriched fe ON f.object_id = fe.object_id + WHERE ($1::text IS NULL OR fth.value = $1) + AND ($2::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $2 || '%')) + ORDER BY fth.timestamp DESC + LIMIT LEAST(COALESCE($3, 100), 1000) + + # FILE_LINKINGS queries + find-linked-files: + kind: postgres-sql + source: chatbot-db + description: Find files linked to a specific file path + parameters: + - name: file_path + type: string + description: The file path to search for (case-insensitive pattern match) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT source, file_path_1, file_path_2, link_type, created_at + FROM file_linkings + WHERE (LOWER(file_path_1) LIKE LOWER('%' || $1 || '%') + OR LOWER(file_path_2) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%')) + ORDER BY created_at DESC + LIMIT LEAST(COALESCE($3, 100), 1000) + + get-file-relationships: + kind: postgres-sql + source: chatbot-db + description: Get all files linked to a specific file path in both directions + parameters: + - name: file_path + type: string + description: Exact file path to find relationships for + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + statement: | + SELECT + CASE + WHEN file_path_1 = $1 THEN file_path_2 + ELSE file_path_1 + END as related_file, + link_type, + source, + created_at + FROM file_linkings + WHERE (file_path_1 = $1 OR file_path_2 = $1) + AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%')) + ORDER BY created_at DESC + + # CHROMIUM.COOKIES queries + count-cookies: + kind: postgres-sql + source: chatbot-db + description: Count browser cookies, optionally filtered by host or source + parameters: + - name: host_pattern + type: string + description: Filter by host_key pattern (case-insensitive, optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + statement: | + SELECT COUNT(*) as cookie_count + FROM chromium.cookies + WHERE ($1::text IS NULL OR LOWER(host_key) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%')) + + search-cookies: + kind: postgres-sql + source: chatbot-db + description: Search browser cookies by host, name, or source + parameters: + - name: host_pattern + type: string + description: Filter by host_key pattern (case-insensitive, optional) + - name: name_pattern + type: string + description: Filter by cookie name pattern (case-insensitive, optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + description: Filter by decryption status (optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT host_key, name, path, value_dec, is_decrypted, is_secure, is_httponly, + expires_utc, source, username, browser, agent_id, project + FROM chromium.cookies + WHERE ($1::text IS NULL OR LOWER(host_key) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(name) LIKE LOWER('%' || $2 || '%')) + AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%')) + AND ($4::boolean IS NULL OR is_decrypted = $4) + ORDER BY expires_utc DESC NULLS LAST + LIMIT LEAST(COALESCE($5, 100), 1000) + + # CHROMIUM.LOGINS queries + count-logins: + kind: postgres-sql + source: chatbot-db + description: Count saved browser credentials, optionally filtered by URL or source + parameters: + - name: url_pattern + type: string + description: Filter by origin_url pattern (case-insensitive, optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + description: Filter by decryption status (optional) + statement: | + SELECT COUNT(*) as login_count + FROM chromium.logins + WHERE ($1::text IS NULL OR LOWER(origin_url) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%')) + AND ($3::boolean IS NULL OR is_decrypted = $3) + + search-logins: + kind: postgres-sql + source: chatbot-db + description: Search browser saved credentials by URL, username, or source + parameters: + - name: url_pattern + type: string + description: Filter by origin_url pattern (case-insensitive, optional) + - name: username_pattern + type: string + description: Filter by username_value pattern (case-insensitive, optional) + - name: source + type: string + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + description: Filter by decryption status (optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT origin_url, username_value, password_value_dec, signon_realm, + is_decrypted, date_created, date_last_used, times_used, + source, username, browser, agent_id, project + FROM chromium.logins + WHERE ($1::text IS NULL OR LOWER(origin_url) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(username_value) LIKE LOWER('%' || $2 || '%')) + AND ($3::text IS NULL OR LOWER(source) LIKE LOWER('%' || $3 || '%')) + AND ($4::boolean IS NULL OR is_decrypted = $4) + ORDER BY date_last_used DESC NULLS LAST + LIMIT LEAST(COALESCE($5, 100), 1000) + + find-lateral-movement-credentials: + kind: postgres-sql + source: chatbot-db + description: Find decrypted credentials that could enable lateral movement to other hosts + parameters: + - name: source + type: string + description: Source/host to find credentials FROM (case-insensitive pattern match, optional) + - name: target_host_pattern + type: string + description: Filter credentials TO specific target hosts (case-insensitive, optional) + - name: limit + type: integer + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT origin_url, username_value, password_value_dec, signon_realm, + date_last_used, times_used, source, username, browser + FROM chromium.logins + WHERE is_decrypted = true + AND password_value_dec IS NOT NULL + AND password_value_dec != '' + AND ($1::text IS NULL OR LOWER(source) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(origin_url) LIKE LOWER('%' || $2 || '%')) + ORDER BY times_used DESC, date_last_used DESC NULLS LAST + LIMIT LEAST(COALESCE($3, 100), 1000) diff --git a/projects/agents/agents/tasks/chatbot.py b/projects/agents/agents/tasks/chatbot.py new file mode 100644 index 0000000..9fcd505 --- /dev/null +++ b/projects/agents/agents/tasks/chatbot.py @@ -0,0 +1,313 @@ +"""Chatbot agent for interactive querying of Nemesis data.""" + +import asyncio +import os +import subprocess +from pathlib import Path +from typing import AsyncGenerator + +import structlog +from agents.base_agent import BaseAgent +from agents.logger import set_agent_metadata +from agents.model_manager import ModelManager +from agents.prompt_manager import PromptManager +from common.db import get_postgres_connection_str +from fastapi import HTTPException +from fastapi.responses import StreamingResponse +from pydantic import BaseModel +from pydantic_ai import Agent +from pydantic_ai.mcp import MCPServerStreamableHTTP +from pydantic_ai.settings import ModelSettings + +logger = structlog.get_logger(__name__) + + +class ChatMessage(BaseModel): + """A single chat message.""" + + role: str # "user" or "assistant" + content: str + + +class ChatbotRequest(BaseModel): + """Request model for chatbot queries.""" + + message: str + history: list[ChatMessage] = [] + use_history: bool = True + temperature: float = 0.7 + + +class ChatbotAgent(BaseAgent): + """Agent for interactive data querying via natural language.""" + + def __init__(self): + super().__init__() + self.prompt_manager = PromptManager(get_postgres_connection_str()) + self.name = "chatbot" + self.description = "Interactive chatbot for querying Nemesis data" + self.agent_type = "llm_based" + self.has_prompt = True + self.llm_temperature = 0.7 # Default, can be overridden per request + + # System prompt - will be saved to DB on first use + self.system_prompt = """You are a technical security analyst assistant for Nemesis, an offensive security data platform. + +Your role is to help users query and analyze security findings, credentials, files, and other artifacts collected during security assessments. + +You have access to a PostgreSQL database with the following tables: +- files_enriched: Processed files with metadata (path, filename, extension, size, magic_type, hashes, etc.) +- enrichments: Detailed analysis results from various enrichment modules (module_name, result_data) +- findings: Security findings categorized by severity and type (finding_name, category, severity, data) +- file_linkings: Relationships between files showing connections (source, file_path_1, file_path_2, link_type) +- chromium.cookies: Browser cookies from Chromium-based browsers (host_key, name, value, expiration) +- chromium.logins: Saved credentials from Chromium browsers (origin_url, username_value, password_value) + +When answering questions: +1. Use precise SQL queries to retrieve relevant data +2. Explain findings in clear, security-focused language +3. Highlight potential lateral movement opportunities when credentials or access is involved +4. Categorize findings by risk level when appropriate +5. Be concise but thorough - prioritize actionable intelligence +6. Always consider the offensive security context +7. When asked about specific hosts/sources, use case-insensitive pattern matching +8. Aggregate and summarize large result sets to provide useful insights + +Query Guidelines: +- Current limit: 1000 rows per query maximum +- Use LIMIT clauses to manage large datasets +- Use COUNT(*) to get totals before retrieving detailed data +- Use GROUP BY to aggregate and summarize when appropriate +- Filter by severity, category, or source to narrow results + +Remember: You are assisting red team operators and penetration testers. Focus on operational value and exploitation opportunities.""" + + self.mcp_process = None + + def get_prompt(self) -> str: + """Get the chatbot prompt from database or use default.""" + try: + prompt_data = self.prompt_manager.get_prompt(self.name) + + if prompt_data: + return prompt_data["prompt"] + else: + logger.info("No prompt found in database, initializing with default", agent_name=self.name) + success = self.prompt_manager.save_prompt(self.name, self.system_prompt, self.description) + if success: + logger.info("Default prompt saved to database", agent_name=self.name) + else: + logger.debug( + "Could not save default prompt to database (likely during startup)", agent_name=self.name + ) + + return self.system_prompt + + except Exception as e: + logger.warning("Error managing prompt, using default", agent_name=self.name, error=str(e)) + return self.system_prompt + + def execute(self, ctx, activity_input: dict) -> dict: + """ + Execute method required by BaseAgent. + + Note: ChatbotAgent is designed for interactive HTTP streaming, + not workflow-based execution. Use the chatbot_stream endpoint instead. + """ + logger.warning("ChatbotAgent.execute called but this agent is designed for HTTP streaming only") + return { + "success": False, + "error": "ChatbotAgent does not support workflow execution. Use /agents/chatbot/stream endpoint instead." + } + + def _get_chatbot_connection_string(self) -> str: + """Get PostgreSQL connection string for chatbot read-only user.""" + chatbot_password = os.getenv("CHATBOT_DB_PASSWORD", "chatbot_pass_change_me") + postgres_host = os.getenv("POSTGRES_HOST", "postgres") + postgres_port = os.getenv("POSTGRES_PORT", "5432") + postgres_db = os.getenv("POSTGRES_DB", "enrichment") + postgres_params = os.getenv("POSTGRES_PARAMETERS", "sslmode=disable") + + return f"postgresql://chatbot_readonly:{chatbot_password}@{postgres_host}:{postgres_port}/{postgres_db}?{postgres_params}" + + async def start_mcp_server(self): + """Start the genai-toolbox MCP server as a subprocess listening on HTTP.""" + # Check if process is already running + if self.mcp_process and self.mcp_process.poll() is None: + logger.debug("MCP server already running") + return + + try: + tools_file = Path(__file__).parent.parent / "mcp" / "tools.yaml" + if not tools_file.exists(): + raise FileNotFoundError(f"tools.yaml not found at {tools_file}") + + # Get database connection string for chatbot readonly user + db_url = self._get_chatbot_connection_string() + + # Start genai-toolbox HTTP server (default port 5000) + logger.info("Starting genai-toolbox MCP HTTP server", tools_file=str(tools_file)) + + self.mcp_process = subprocess.Popen( + ["genai-toolbox", "--tools-file", str(tools_file)], + env={**os.environ, "DATABASE_URL": db_url}, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # Give it a moment to start + await asyncio.sleep(2) + + if self.mcp_process.poll() is not None: + stderr = self.mcp_process.stderr.read().decode() if self.mcp_process.stderr else "" + # If it failed due to address in use, that's actually okay + if "address already in use" in stderr.lower(): + logger.info("MCP server already running (address in use)") + self.mcp_process = None + return + raise RuntimeError(f"MCP server failed to start: {stderr}") + + logger.info("MCP HTTP server started successfully on http://127.0.0.1:5000/mcp") + + except Exception as e: + logger.error("Failed to start MCP server", error=str(e)) + raise + + async def stop_mcp_server(self): + """Stop the MCP server subprocess.""" + if self.mcp_process and self.mcp_process.poll() is None: + logger.info("Stopping MCP server") + self.mcp_process.terminate() + try: + self.mcp_process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning("MCP server didn't stop gracefully, killing") + self.mcp_process.kill() + self.mcp_process = None + + async def stream_chat_response(self, request: ChatbotRequest) -> AsyncGenerator[str, None]: + """ + Stream chatbot responses token-by-token. + + Args: + request: ChatbotRequest with message, history, and settings + + Yields: + Chunks of the response as they're generated + """ + model = ModelManager.get_model() + + if not model: + logger.warning("No model available from ModelManager") + raise HTTPException(status_code=503, detail="AI model not available") + + try: + # Set metadata for Phoenix tracing + set_agent_metadata( + agent_name="chatbot", + message_length=len(request.message), + has_history=len(request.history) > 0 if request.use_history else False, + tags=["chatbot", "interactive_query"], + ) + + # Get current prompt from database + current_prompt = self.get_prompt() + + # Build conversation history if enabled + conversation = "" + if request.use_history and request.history: + for msg in request.history: + role_label = "User" if msg.role == "user" else "Assistant" + conversation += f"{role_label}: {msg.content}\n\n" + + # Add current message + conversation += f"User: {request.message}\n\nAssistant:" + + # Connect to MCP HTTP server (genai-toolbox running on http://127.0.0.1:5000/mcp) + mcp_server = MCPServerStreamableHTTP(url='http://127.0.0.1:5000/mcp') + + # Create agent with MCP tools + agent = Agent( + model=model, + system_prompt=current_prompt, + toolsets=[mcp_server], + instrument=ModelManager.is_instrumentation_enabled(), + retries=2, + model_settings=ModelSettings(temperature=request.temperature), + ) + + logger.info("Starting chatbot stream", message=request.message, temperature=request.temperature) + + # When tools are involved, streaming doesn't work as expected + # Get the complete result and send it + result = await agent.run(conversation) + + final_text = result.data if hasattr(result, 'data') else str(result) + logger.info(f"Got complete response, {len(final_text)} chars") + + # Send the complete response + if final_text: + yield final_text + else: + logger.warning("No text in final result") + + # Log tool calls + tool_calls = [] + if hasattr(result, 'all_messages'): + for msg in result.all_messages(): + if hasattr(msg, 'parts'): + for part in msg.parts: + if hasattr(part, 'tool_name'): + tool_calls.append({ + 'tool': part.tool_name, + 'args': getattr(part, 'args', {}) + }) + + if tool_calls: + logger.info("MCP tools called", tool_calls=tool_calls, count=len(tool_calls)) + else: + logger.warning("No MCP tools were called by the LLM") + + # Log completion metrics + logger.info( + "Chatbot response completed", + total_tokens=result.usage().total_tokens if hasattr(result, "usage") else None, + ) + + except Exception as e: + logger.error("Chatbot streaming failed", error=str(e)) + yield f"\n\n[Error: {str(e)}]" + + +# Global chatbot agent instance +_chatbot_agent: ChatbotAgent | None = None + + +def get_chatbot_agent() -> ChatbotAgent: + """Get or create the global chatbot agent instance.""" + global _chatbot_agent + if _chatbot_agent is None: + _chatbot_agent = ChatbotAgent() + return _chatbot_agent + + +async def chatbot_stream(request: ChatbotRequest) -> StreamingResponse: + """ + FastAPI endpoint handler for streaming chatbot responses. + + Args: + request: ChatbotRequest with message and settings + + Returns: + StreamingResponse with text/event-stream content + """ + agent = get_chatbot_agent() + + # Ensure MCP HTTP server is running + await agent.start_mcp_server() + + # Stream the response (connects to MCP server via HTTP) + return StreamingResponse( + agent.stream_chat_response(request), + media_type="text/plain", + ) diff --git a/projects/frontend/src/App.jsx b/projects/frontend/src/App.jsx index df04137..4350c69 100644 --- a/projects/frontend/src/App.jsx +++ b/projects/frontend/src/App.jsx @@ -12,6 +12,7 @@ import { HelpCircle, Key, LayoutDashboard, + MessageSquare, Search, Settings, Siren, @@ -41,6 +42,7 @@ import ThemeToggle from './components/ThemeToggle'; import YaraRulesManager from './components/Yara/YaraManager'; import Containers from './components/Containers/Containers'; import AgentsPage from './components/Agents/AgentsPage'; +import ChatbotPage from './components/Chatbot/ChatbotPage'; import FileBrowser from './components/FileBrowser/FileBrowser'; import Chromium from './components/Chromium/Chromium'; import Dpapi from './components/Dpapi/Dpapi'; @@ -197,9 +199,13 @@ const Sidebar = ({ onCollapse }) => { { id: 'reporting', label: 'Reporting', icon: BarChart2, path: '/reporting' } ]; - // Add Agents tab if LiteLLM is available + // Add Chatbot and Agents tabs if LiteLLM is available const navigationItems = litellmAvailable - ? [...baseNavigationItems, { id: 'agents', label: 'Agents', icon: Bot, path: '/agents' }] + ? [ + ...baseNavigationItems, + { id: 'chatbot', label: 'Chatbot', icon: MessageSquare, path: '/chatbot' }, + { id: 'agents', label: 'Agents', icon: Bot, path: '/agents' } + ] : baseNavigationItems; const utilityItems = [ @@ -369,6 +375,7 @@ const App = () => { } /> } /> } /> + } /> } /> } /> } /> diff --git a/projects/frontend/src/components/Chatbot/ChatbotPage.jsx b/projects/frontend/src/components/Chatbot/ChatbotPage.jsx new file mode 100644 index 0000000..8273e19 --- /dev/null +++ b/projects/frontend/src/components/Chatbot/ChatbotPage.jsx @@ -0,0 +1,291 @@ +import React, { useState, useEffect, useRef } from 'react'; +import { Bot, Send, Trash2, Settings as SettingsIcon, AlertCircle } from 'lucide-react'; +import ExampleQueries from './ExampleQueries'; +import MessageBubble from './MessageBubble'; +import QueryModal from './QueryModal'; + +const ChatbotPage = () => { + const [messages, setMessages] = useState([]); + const [currentMessage, setCurrentMessage] = useState(''); + const [isStreaming, setIsStreaming] = useState(false); + const [useHistory, setUseHistory] = useState(true); + const [showQueries, setShowQueries] = useState(false); + const [temperature, setTemperature] = useState(0.7); + const [queries, setQueries] = useState([]); + const [error, setError] = useState(null); + const [showSettings, setShowSettings] = useState(false); + + const messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + + // Auto-scroll to bottom when messages change + const scrollToBottom = () => { + messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }); + }; + + useEffect(() => { + scrollToBottom(); + }, [messages]); + + const sendMessage = async (messageText) => { + if (!messageText.trim() || isStreaming) return; + + const userMessage = { role: 'user', content: messageText.trim() }; + setMessages(prev => [...prev, userMessage]); + setCurrentMessage(''); + setIsStreaming(true); + setError(null); + + // Create abort controller for this request + abortControllerRef.current = new AbortController(); + + try { + const response = await fetch('/api/chatbot/stream', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + message: messageText.trim(), + history: useHistory ? messages : [], + use_history: useHistory, + temperature: temperature + }), + signal: abortControllerRef.current.signal + }); + + if (!response.ok) { + throw new Error(`HTTP ${response.status}: ${response.statusText}`); + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let assistantMessage = ''; + + // Add empty assistant message that we'll update + setMessages(prev => [...prev, { role: 'assistant', content: '' }]); + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + const chunk = decoder.decode(value, { stream: true }); + assistantMessage += chunk; + + // Update the last message (assistant's response) in real-time + setMessages(prev => { + const updated = [...prev]; + updated[updated.length - 1] = { + role: 'assistant', + content: assistantMessage + }; + return updated; + }); + } + + // TODO: Extract SQL queries from response if showQueries is enabled + // For now, queries would need to be returned in a structured format from backend + + } catch (err) { + if (err.name === 'AbortError') { + console.log('Request aborted'); + } else { + console.error('Streaming error:', err); + setError(err.message || 'Failed to get response from chatbot'); + + // Remove the empty assistant message if there was an error + setMessages(prev => prev.filter(msg => msg.content !== '')); + } + } finally { + setIsStreaming(false); + abortControllerRef.current = null; + } + }; + + const handleSubmit = (e) => { + e.preventDefault(); + sendMessage(currentMessage); + }; + + const handleExampleClick = (example) => { + setCurrentMessage(example); + }; + + const clearHistory = () => { + if (window.confirm('Clear all conversation history?')) { + setMessages([]); + setQueries([]); + setError(null); + } + }; + + const handleKeyPress = (e) => { + if (e.key === 'Enter' && !e.shiftKey) { + e.preventDefault(); + handleSubmit(e); + } + }; + + return ( +
+ {/* Header */} +
+
+ +
+

Chatbot

+

+ Query Nemesis data with natural language +

+
+
+ +
+ {/* Settings Toggle */} + + + {/* Clear History */} + +
+
+ + {/* Settings Panel */} + {showSettings && ( +
+

Settings

+ +
+ {/* Temperature Control */} +
+ + setTemperature(parseFloat(e.target.value))} + className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer dark:bg-gray-700" + /> +
+ + {/* Use History Toggle */} +
+ Use Conversation History + +
+ + {/* Show Queries Toggle */} +
+ Show SQL Queries + +
+
+
+ )} + + {/* Error Display */} + {error && ( +
+ +
+

{error}

+
+
+ )} + + {/* Example Queries (show when no messages) */} + {messages.length === 0 && ( + + )} + + {/* Messages Container */} +
+ {messages.map((msg, idx) => ( + + ))} +
+
+ + {/* Query Modal */} + {showQueries && queries.length > 0 && ( + + )} + + {/* Input Form */} +
+
+