diff --git a/compose.yaml b/compose.yaml index 9b051e8..f4d84ed 100644 --- a/compose.yaml +++ b/compose.yaml @@ -1049,6 +1049,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_RESULTS=${MCP_MAX_RESULTS:-1000} logging: *logging-config depends_on: postgres: { condition: service_healthy } diff --git a/env.example b/env.example index 1928b65..bd1a933 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 results returned by chatbot tools +# Helps prevent expensive queries and context window issues. +# Example: +# MCP_MAX_RESULTS=500 +MCP_MAX_RESULTS=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..803f1a3 --- /dev/null +++ b/projects/agents/agents/mcp/tools.yaml @@ -0,0 +1,506 @@ +# 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 + required: false + description: Filter by project name (optional) + - name: agent_id + type: string + required: false + description: Filter by agent ID (optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: extension + type: string + required: false + 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 + required: false + description: Search pattern for filename (case-insensitive, optional) + - name: path_pattern + type: string + required: false + description: Search pattern for file path (case-insensitive, optional) + - name: extension + type: string + required: false + description: Filter by file extension (optional) + - name: project + type: string + required: false + description: Filter by project name (optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + required: false + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT object_id::text, file_name, path, extension, size, magic_type, mime_type, + source, agent_id, project, timestamp, originating_object_id::text + 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 object_id::text, agent_id, source, project, timestamp, path, file_name, + extension, size, magic_type, mime_type, originating_object_id::text + 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.object_id::text, 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 + required: false + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT e.object_id::text, 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 + get-unique-findings-categories: + kind: postgres-sql + source: chatbot-db + description: List all unique finding categories with their counts and severity statistics + statement: | + SELECT category, COUNT(*) as count, + AVG(severity) as avg_severity, + MAX(severity) as max_severity + FROM findings + GROUP BY category + ORDER BY count DESC + + count-findings: + kind: postgres-sql + source: chatbot-db + description: Count findings, optionally filtered by severity, category, or source + parameters: + - name: min_severity + type: integer + required: false + description: Minimum severity level (0-10, optional) + - name: category + type: string + required: false + description: Filter by finding category (optional) + - name: source + type: string + required: false + 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 $2 = '' OR f.category = $2) + AND ($3::text IS NULL OR $3 = '' 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 + required: false + description: Minimum severity level (0-10, optional) + - name: category + type: string + required: false + description: Filter by finding category (optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + required: false + 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-file-findings: + kind: postgres-sql + source: chatbot-db + description: Get all findings for a specific file by object_id + parameters: + - name: object_id + type: string + description: The UUID of the file + 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 f.object_id = $1::uuid + ORDER BY f.severity DESC, f.created_at DESC + + get-findings-by-category: + kind: postgres-sql + source: chatbot-db + description: Get aggregated count of findings grouped by category + parameters: + - name: source + type: string + required: false + 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 + required: false + description: Filter by triage value (e.g., true_positive, false_positive, optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + required: false + 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) + + search-credential-findings-by-host: + kind: postgres-sql + source: chatbot-db + description: Search for credential findings related to a specific hostname or system. Use this to find credentials that may provide access to a target system. + parameters: + - name: hostname + type: string + description: Target hostname or system name to search for in credential findings (case-insensitive partial match) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + required: false + 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 f.category = 'credential' + AND LOWER(f.data::text) LIKE LOWER('%' || $1 || '%') + AND ($2::text IS NULL OR LOWER(fe.source) LIKE LOWER('%' || $2 || '%')) + ORDER BY f.severity DESC, f.created_at DESC + LIMIT LEAST(COALESCE($3, 100), 1000) + + search-logins-by-host: + kind: postgres-sql + source: chatbot-db + description: Search for decrypted browser credentials related to a specific hostname or URL. Use this to find saved browser passwords for accessing a target system. + parameters: + - name: hostname + type: string + description: Target hostname or URL to search for in saved browser credentials (case-insensitive partial match) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: limit + type: integer + required: false + description: Maximum number of results (default 100, max 1000) + statement: | + SELECT origin_url, username_value, password_value_dec, signon_realm, + date_created, date_last_used, times_used, + source, username, browser, agent_id, project + FROM chromium.logins + WHERE is_decrypted = true + AND (LOWER(origin_url) LIKE LOWER('%' || $1 || '%') + OR LOWER(signon_realm) LIKE LOWER('%' || $1 || '%')) + AND ($2::text IS NULL OR LOWER(source) LIKE LOWER('%' || $2 || '%')) + ORDER BY date_last_used DESC NULLS LAST, times_used DESC + LIMIT LEAST(COALESCE($3, 100), 1000) + + # 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 + required: false + description: Filter by host_key pattern (case-insensitive, optional) + - name: source + type: string + required: false + 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 + required: false + description: Filter by host_key pattern (case-insensitive, optional) + - name: name_pattern + type: string + required: false + description: Filter by cookie name pattern (case-insensitive, optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + required: false + description: Filter by decryption status (optional) + - name: limit + type: integer + required: false + 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, originating_object_id::text + 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 + required: false + description: Filter by origin_url pattern (case-insensitive, optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + required: false + 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 + required: false + description: Filter by origin_url pattern (case-insensitive, optional) + - name: username_pattern + type: string + required: false + description: Filter by username_value pattern (case-insensitive, optional) + - name: source + type: string + required: false + description: Filter by source (case-insensitive pattern match, optional) + - name: is_decrypted + type: boolean + required: false + description: Filter by decryption status (optional) + - name: limit + type: integer + required: false + 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, originating_object_id::text + 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) + + # PLAINTEXT CONTENT searches + search-document-content: + kind: postgres-sql + source: chatbot-db + description: Full-text search through plaintext content chunks of files. Returns the first matching chunk from each file. + parameters: + - name: search_query + type: string + description: Text to search for in document content (full-text search) + - name: path_pattern + type: string + required: false + description: Filter by file path pattern using LIKE (e.g., '%folder%', optional) + - name: agent_pattern + type: string + required: false + description: Filter by agent ID pattern using LIKE (optional) + - name: project_name + type: string + required: false + description: Filter by exact project name (optional) + - name: start_date + type: string + required: false + description: Filter files from this date onwards in ISO format (e.g., '2024-01-01', optional) + - name: end_date + type: string + required: false + description: Filter files up to this date in ISO format (e.g., '2024-12-31', optional) + - name: max_results + type: integer + required: false + description: Maximum number of results (default 100, max 1000) + - name: source_pattern + type: string + required: false + description: Filter by source pattern using LIKE (optional) + statement: | + SELECT object_id::text, chunk_number, content, file_name, path, extension, + project, agent_id, source, timestamp + FROM public.search_documents( + $1, + $2, + $3, + $4, + $5::timestamp with time zone, + $6::timestamp with time zone, + COALESCE($7, 100), + $8 + ) diff --git a/projects/agents/agents/tasks/chatbot.py b/projects/agents/agents/tasks/chatbot.py new file mode 100644 index 0000000..9a9afcb --- /dev/null +++ b/projects/agents/agents/tasks/chatbot.py @@ -0,0 +1,355 @@ +"""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 + + # Get max rows from environment + max_results = int(os.getenv("MCP_MAX_RESULTS", "1000")) + + # System prompt - will be saved to DB on first use + self.system_prompt = f"""You are a data query assistant for Nemesis, an offensive security data platform. + +Your role is to retrieve and report data from the database. Do NOT provide recommendations, analysis, or suggestions - only report the requested data. You have access to MCP tools to query the Nemesis PostgreSQL database. + +When answering questions: +1. Query the database using the appropriate tools +2. Report ONLY the data retrieved unless explicitly instructed otherwise - i.e., no analysis or recommendations unless a user explicitly asks for it +3. Present results clearly and concisely +4. For large result sets, summarize counts and key details +5. Use case-insensitive pattern matching for host/source filters +6. Be brief - users want facts, not explanations unless they explicitly request them +7. Don't return the "project" field to users + +Query Guidelines: +- Use a `limit` of {max_results} for maximum results +- Filter by severity, category, or source to narrow results +- Use search-document-content when users ask to search "for" or "containing" specific text but otherwise restrict your usage of "search-document-content" since it can return a lot of results +- a `originating_object_id` field points to the `object_id` the finding/file originated from + +Searching for Credentials to Access Systems: +When users ask about accessing a specific system or finding credentials for a hostname, follow this order: +1. First, use search-credential-findings-by-host to find credential findings related to the target hostname +2. Second, use search-logins-by-host to find decrypted browser credentials for the target hostname +3. Only as a last resort, if the above return no results, use search-document-content with the hostname to search file contents +This order ensures you check the most relevant credential sources first before falling back to broader document searches. +""" + + 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=5, # Increased from 2 to handle transient MCP tool failures + 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) + + # Log tool calls and their results with full details + tool_calls = [] + tool_errors = [] + if hasattr(result, 'all_messages'): + for msg_idx, msg in enumerate(result.all_messages()): + logger.debug(f"Message {msg_idx}: type={type(msg).__name__}, role={getattr(msg, 'role', 'unknown')}") + if hasattr(msg, 'parts'): + for part_idx, part in enumerate(msg.parts): + part_type = type(part).__name__ + logger.debug(f" Part {part_idx}: type={part_type}") + + # Log tool calls (requests) + if hasattr(part, 'tool_name'): + tool_info = { + 'tool': part.tool_name, + 'args': getattr(part, 'args', {}) + } + logger.info(f"TOOL CALL: {part.tool_name}", args=tool_info['args']) + tool_calls.append(tool_info) + + # Log tool returns (responses) + if hasattr(part, 'tool_name') and hasattr(part, 'content'): + full_content = str(part.content) + logger.debug( + f"TOOL RESPONSE: {part.tool_name}", + content_length=len(full_content), + full_content=full_content # Log FULL content for debugging + ) + + # Log errors + if hasattr(part, 'error'): + error_info = { + 'tool': getattr(part, 'tool_name', 'unknown'), + 'error': str(part.error) + } + logger.error("TOOL ERROR", error_info=error_info) + tool_errors.append(error_info) + + if tool_calls: + logger.info("MCP tools called", count=len(tool_calls), tools=[t['tool'] for t in tool_calls]) + else: + logger.warning("No MCP tools were called by the LLM") + + if tool_errors: + logger.error("Tool errors occurred", error_count=len(tool_errors)) + + # Extract just the text output from the result + if hasattr(result, 'data'): + final_text = str(result.data) + elif hasattr(result, 'output'): + final_text = str(result.output) + else: + final_text = str(result) + + logger.info(f"Got complete response, {len(final_text)} chars") + + # Debug: Log final response to check for UUID corruption + if 'object_id' in final_text.lower() or 'uuid' in final_text.lower(): + logger.warning( + "FINAL RESPONSE contains object_id/UUID", + final_response=final_text[:2000] # Log first 2000 chars + ) + + # Send the complete response + if final_text: + yield final_text + else: + logger.warning("No text in final result") + + # 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/file_enrichment/pyproject.toml b/projects/file_enrichment/pyproject.toml index 87d4517..c45f542 100644 --- a/projects/file_enrichment/pyproject.toml +++ b/projects/file_enrichment/pyproject.toml @@ -43,6 +43,7 @@ pillow = "^11.3.0" opentelemetry-api = "^1.38.0" opentelemetry-sdk = "^1.38.0" opentelemetry-exporter-otlp-proto-grpc = "^1.38.0" +presidio-analyzer = "^2.2.360" [tool.poetry.group.dev.dependencies] ruff = "^0.9.2" 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..a85fcd7 --- /dev/null +++ b/projects/frontend/src/components/Chatbot/ChatbotPage.jsx @@ -0,0 +1,520 @@ +import { AlertCircle, Bot, Send, Settings as SettingsIcon, Trash2, X } from 'lucide-react'; +import { useEffect, useRef, useState } from 'react'; +import ExampleQueries from './ExampleQueries'; +import MessageBubble from './MessageBubble'; + +const MESSAGES_STORAGE_KEY = 'chatbot_messages'; +const TOKEN_WARNING_THRESHOLD = 50000; + +// Rough token estimation: ~4 characters per token +const estimateTokens = (text) => { + return Math.ceil(text.length / 4); +}; + +const ChatbotPage = () => { + const [messages, setMessages] = useState(() => { + // Load messages from localStorage on mount + try { + const saved = localStorage.getItem(MESSAGES_STORAGE_KEY); + return saved ? JSON.parse(saved) : []; + } catch { + return []; + } + }); + const [currentMessage, setCurrentMessage] = useState(''); + const [isStreaming, setIsStreaming] = useState(false); + const [useHistory, setUseHistory] = useState(true); + const [temperature, setTemperature] = useState(0.7); + const [error, setError] = useState(null); + const [showSettings, setShowSettings] = useState(false); + const [systemPrompt, setSystemPrompt] = useState(''); + const [originalPrompt, setOriginalPrompt] = useState(''); + const [savingPrompt, setSavingPrompt] = useState(false); + const [promptError, setPromptError] = useState(null); + const [spendData, setSpendData] = useState(null); + const [showTokenWarning, setShowTokenWarning] = useState(false); + const [estimatedTokens, setEstimatedTokens] = useState(0); + + const messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + const textareaRef = useRef(null); + + // Auto-scroll to bottom when messages change + const scrollToBottom = () => { + messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }); + }; + + useEffect(() => { + scrollToBottom(); + }, [messages]); + + // Save messages to localStorage whenever they change + useEffect(() => { + try { + localStorage.setItem(MESSAGES_STORAGE_KEY, JSON.stringify(messages)); + } catch (err) { + console.error('Failed to save messages to localStorage:', err); + } + }, [messages]); + + // Refocus input after response completes + useEffect(() => { + if (!isStreaming && messages.length > 0) { + // Small delay to ensure DOM has updated + setTimeout(() => { + textareaRef.current?.focus(); + }, 100); + } + }, [isStreaming, messages.length]); + + // Fetch system prompt when settings are opened + useEffect(() => { + if (showSettings && !systemPrompt) { + fetchSystemPrompt(); + } + }, [showSettings]); + + // Fetch spend data on mount + useEffect(() => { + fetchSpendData(); + }, []); + + // Calculate token usage and show warning if needed + useEffect(() => { + const totalText = messages.map(m => m.content).join(' '); + const tokens = estimateTokens(totalText); + setEstimatedTokens(tokens); + + if (tokens > TOKEN_WARNING_THRESHOLD && !showTokenWarning) { + setShowTokenWarning(true); + } + }, [messages]); + + const fetchSpendData = async () => { + try { + const response = await fetch('/api/agents/spend-data'); + if (!response.ok) throw new Error('Network response error'); + const result = await response.json(); + + setSpendData({ + spend: result.total_spend, + total_tokens: result.total_tokens + }); + } catch (err) { + console.error('Error fetching spend data:', err); + // Don't set error state - gracefully degrade + } + }; + + const fetchSystemPrompt = async () => { + try { + const query = { + query: ` + query GetChatbotPrompt { + agent_prompts_by_pk(name: "chatbot") { + name + prompt + description + } + } + ` + }; + + const response = await fetch('/hasura/v1/graphql', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-hasura-admin-secret': window.ENV.HASURA_ADMIN_SECRET, + }, + body: JSON.stringify(query) + }); + + if (!response.ok) throw new Error('Network response error'); + const result = await response.json(); + if (result.errors) throw new Error(result.errors[0].message); + + const prompt = result.data.agent_prompts_by_pk?.prompt || ''; + setSystemPrompt(prompt); + setOriginalPrompt(prompt); + } catch (err) { + console.error('Error fetching system prompt:', err); + setPromptError('Failed to load system prompt'); + } + }; + + const saveSystemPrompt = async () => { + if (systemPrompt === originalPrompt) return; + + setSavingPrompt(true); + setPromptError(null); + + try { + const mutation = { + query: ` + mutation UpsertChatbotPrompt($prompt: String!) { + insert_agent_prompts_one( + object: { + name: "chatbot", + prompt: $prompt, + description: "Interactive chatbot for querying Nemesis data" + }, + on_conflict: { + constraint: agent_prompts_pkey, + update_columns: [prompt] + } + ) { + name + prompt + } + } + `, + variables: { prompt: systemPrompt } + }; + + const response = await fetch('/hasura/v1/graphql', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-hasura-admin-secret': window.ENV.HASURA_ADMIN_SECRET, + }, + body: JSON.stringify(mutation) + }); + + if (!response.ok) throw new Error('Network response error'); + const result = await response.json(); + if (result.errors) throw new Error(result.errors[0].message); + + setOriginalPrompt(systemPrompt); + alert('System prompt saved successfully!'); + } catch (err) { + console.error('Error saving system prompt:', err); + setPromptError('Failed to save system prompt'); + } finally { + setSavingPrompt(false); + } + }; + + const cancelPromptEdit = () => { + setSystemPrompt(originalPrompt); + setPromptError(null); + }; + + 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; + }); + } + + } 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; + // Refresh spend data after response completes + fetchSpendData(); + } + }; + + const handleSubmit = (e) => { + e.preventDefault(); + sendMessage(currentMessage); + }; + + const handleExampleClick = (example) => { + setCurrentMessage(example); + }; + + const clearHistory = () => { + if (window.confirm('Clear all conversation history?')) { + setMessages([]); + setError(null); + setShowTokenWarning(false); + // Clear from localStorage as well + try { + localStorage.removeItem(MESSAGES_STORAGE_KEY); + } catch (err) { + console.error('Failed to clear messages from localStorage:', err); + } + } + }; + + const handleKeyPress = (e) => { + if (e.key === 'Enter' && !e.shiftKey) { + e.preventDefault(); + handleSubmit(e); + } + }; + + return ( +
+ {/* Header */} +
+
+ +
+

Chatbot

+

+ Query Nemesis data with natural language +

+
+
+ + {/* LLM Usage Stats */} + {spendData && ( +
+
+
+ ${spendData.spend ? spendData.spend.toFixed(4) : '0.0000'} +
+
Total Spend
+
+
+
+ {spendData.total_tokens ? spendData.total_tokens.toLocaleString() : '0'} +
+
Total Tokens
+
+
+ )} + +
+ {/* 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 + +
+ + {/* System Prompt Editor */} +
+

System Prompt

+ + {promptError && ( +
+ {promptError} +
+ )} + +