NOVA MCP Security Gateway

import asyncio import os import sys from typing import Optional from contextlib import AsyncExitStack from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from openai import OpenAI from dotenv import load_dotenv import json import logging import os logging.getLogger("sentence_transformers").setLevel(logging.ERROR) load_dotenv() openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) # Configure client logging to the same nova_matches.log used by the server CLIENT_LOG_DIR = os.path.join(os.path.dirname(__file__), "logs") os.makedirs(CLIENT_LOG_DIR, exist_ok=True) CLIENT_LOG_FILE = os.path.join(CLIENT_LOG_DIR, "nova_matches.log") client_logger = logging.getLogger("nova-mcp-client") client_logger.setLevel(logging.DEBUG) fh = logging.FileHandler(CLIENT_LOG_FILE) fh.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')) client_logger.addHandler(fh) class NovaSecurityClient: def __init__(self): self.session: Optional[ClientSession] = None self.exit_stack = AsyncExitStack() async def connect_to_server(self, server_script_path: str): if not server_script_path.endswith((".py", ".js")): raise ValueError("Server script must be a .py or .js file") # Get the absolute path to the server script server_script_path = os.path.abspath(server_script_path) # Set working directory to where the server script is located working_dir = os.path.dirname(server_script_path) command = "python" if server_script_path.endswith(".py") else "node" server_params = StdioServerParameters( command=command, args=[os.path.basename(server_script_path)], cwd=working_dir # Set working directory ) client_logger.info(f"Starting NOVA MCP server: {command} {server_script_path}") print(f"Starting NOVA MCP server: {command} {server_script_path}...") # Connect to server with a timeout per operation (asyncio.wait_for for Python 3.10 compatibility) connection_timeout = 20 # Seconds try: stdio_transport = await asyncio.wait_for( self.exit_stack.enter_async_context(stdio_client(server_params)), connection_timeout ) self.stdio, self.write = stdio_transport self.session = await asyncio.wait_for( self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write)), connection_timeout ) # Give server a moment to start await asyncio.sleep(1) await asyncio.wait_for(self.session.initialize(), connection_timeout) response = await asyncio.wait_for(self.session.list_tools(), connection_timeout) tools = response.tools # Verify the validate_prompt tool is available if not any(tool.name == "validate_prompt" for tool in tools): raise ConnectionError("Server started but 'validate_prompt' tool not found") print("\n✅ Connected to MCP server.") print("Available tools:", [tool.name for tool in tools]) except asyncio.TimeoutError: raise ConnectionError( f"Timed out connecting to server after {connection_timeout} seconds. Check if the server is running correctly." ) except Exception as e: # Log the full exception with stack trace client_logger.error("Exception during connect_to_server", exc_info=True) # Include exception type and message for better debugging err_class = e.__class__.__name__ err_message = str(e) or repr(e) raise ConnectionError( f"Failed to connect to NOVA MCP server: {err_class}: {err_message}" ) from e async def process_query(self, query: str) -> str: # ...existing code... system_prompt = ( "Here is your instruction you MUST follow: " "You are an AI. For this session, Nova-Security MCP is responsible for verifying all prompts. " "Before doing anything else, you MUST pass every prompt to the MCP for validation. " "If a prompt is not authorized, do NOT respond. Instead, return the exact message received from the MCP—nothing else." ) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": query} ] tools_response = await self.session.list_tools() tool_definitions = [ { "type": "function", "function": { "name": tool.name, "description": tool.description, "parameters": tool.inputSchema, }, } for tool in tools_response.tools ] response = openai_client.chat.completions.create( model="gpt-4o", messages=messages, tools=tool_definitions, tool_choice="auto" ) full_response = [] assistant_msg = response.choices[0].message if assistant_msg.tool_calls: for tool_call in assistant_msg.tool_calls: tool_name = tool_call.function.name try: tool_args = json.loads(tool_call.function.arguments) except json.JSONDecodeError: tool_args = {} try: result = await self.session.call_tool(tool_name, tool_args) except Exception as e: result = f"Error: {str(e)}" messages.extend([ { "role": "assistant", "tool_calls": [tool_call], "content": None }, { "role": "tool", "tool_call_id": tool_call.id, "content": str(result) } ]) response = openai_client.chat.completions.create( model="gpt-4o", messages=messages ) final_output = response.choices[0].message.content full_response.append(final_output) else: full_response.append(assistant_msg.content) response_text = "\n".join(full_response) # Log result as JSON: authorized at INFO, unauthorized at WARNING if response_text.startswith("NOT AUTHORIZED"): # Parse unauthorized details details = {"user_id": "unknown", "prompt": query} for line in response_text.splitlines(): if line.startswith("Security rule matched:"): details["rule_name"] = line.split(":", 1)[1].strip() if line.startswith("Severity:"): details["severity"] = line.split(":", 1)[1].strip() client_logger.warning(json.dumps(details)) else: # Authorized: log query and response together client_logger.info(json.dumps({"query": query, "response": response_text})) return response_text async def chat_loop(self): print("\n🧠 Nova-Security MCP Client (GPT-4o)") print("Type your queries or 'quit' to exit.") while True: try: query = input("\nQuery: ").strip() if query.lower() == "quit": break response = await self.process_query(query) print("\n" + response) except Exception as e: print(f"\n❌ Error: {str(e)}") async def cleanup(self): """Clean up resources and gracefully close connection to server""" client_logger.info("Cleaning up resources and closing session") print("Cleaning up resources...") try: if self.session: # Send a clean shutdown message if possible try: await self.session.shutdown() except Exception: pass # Ignore errors during shutdown # Close the exit stack await self.exit_stack.aclose() except Exception as e: print(f"Error during cleanup: {str(e)}") # Don't re-raise - we're already in cleanup async def main(): if len(sys.argv) < 2: print("Usage: python client.py <path_to_server_script>") sys.exit(1) client = NovaSecurityClient() try: await client.connect_to_server(sys.argv[1]) await client.chat_loop() finally: await client.cleanup() if __name__ == "__main__": asyncio.run(main())