#! /usr/bin/python3 # This script connects the MCP AI agent to Kali Linux terminal and API Server. # Inspired by https://github.com/whit3rabbit0/project_astro import argparse import json import logging import os import re import shlex import subprocess import sys import traceback import threading from typing import Dict, Any, Union, List from flask import Flask, request, jsonify # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.StreamHandler(sys.stdout) ] ) logger = logging.getLogger(__name__) # Configuration API_PORT = int(os.environ.get("API_PORT", 5000)) DEBUG_MODE = os.environ.get("DEBUG_MODE", "0").lower() in ("1", "true", "yes", "y") COMMAND_TIMEOUT = 180 # 5 minutes default timeout app = Flask(__name__) class CommandExecutor: """Class to handle command execution with better timeout management""" def __init__(self, command: Union[str, List[str]], timeout: int = COMMAND_TIMEOUT): self.command = command self.timeout = timeout # We only use shell mode if the input is a plain string. # Lists will run safely without a shell wrapper. self.use_shell = isinstance(command, str) self.process = None self.stdout_data = "" self.stderr_data = "" self.stdout_thread = None self.stderr_thread = None self.return_code = None self.timed_out = False def _read_stdout(self): """Thread function to continuously read stdout""" for line in iter(self.process.stdout.readline, ''): self.stdout_data += line def _read_stderr(self): """Thread function to continuously read stderr""" for line in iter(self.process.stderr.readline, ''): self.stderr_data += line def execute(self) -> Dict[str, Any]: """Execute the command and handle timeout gracefully""" logger.info(f"Executing command: {self.command}") # Fix: Parse strings safely into an execution array, or preserve arrays directly if isinstance(self.command, str): final_args = shlex.split(self.command) elif isinstance(self.command, list): final_args = self.command else: raise ValueError(f"CommandExecutor expects a string or list, but got {type(self.command).__name__}") try: # Fix: Pass final_args instead of the unparsed command object self.process = subprocess.Popen( final_args, shell=False, # Setting shell=False is significantly more secure for variable inputs stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1 # Line buffered ) # Start threads to read output continuously self.stdout_thread = threading.Thread(target=self._read_stdout) self.stderr_thread = threading.Thread(target=self._read_stderr) self.stdout_thread.daemon = True self.stderr_thread.daemon = True self.stdout_thread.start() self.stderr_thread.start() # Wait for the process to complete or timeout try: self.return_code = self.process.wait(timeout=self.timeout) # Process completed, join the threads self.stdout_thread.join() self.stderr_thread.join() except subprocess.TimeoutExpired: # Process timed out but we might have partial results self.timed_out = True logger.warning(f"Command timed out after {self.timeout} seconds. Terminating process.") # Try to terminate gracefully first self.process.terminate() try: self.process.wait(timeout=5) # Give it 5 seconds to terminate except subprocess.TimeoutExpired: # Force kill if it doesn't terminate logger.warning("Process not responding to termination. Killing.") self.process.kill() # Update final output self.return_code = -1 # Always consider it a success if we have output, even with timeout success = True if self.timed_out and (self.stdout_data or self.stderr_data) else (self.return_code == 0) return { "stdout": self.stdout_data, "stderr": self.stderr_data, "return_code": self.return_code, "success": success, "timed_out": self.timed_out, "partial_results": self.timed_out and (self.stdout_data or self.stderr_data) } except Exception as e: logger.error(f"Error executing command: {str(e)}") logger.error(traceback.format_exc()) return { "stdout": self.stdout_data, "stderr": f"Error executing command: {str(e)}\n{self.stderr_data}", "return_code": -1, "success": False, "timed_out": False, "partial_results": bool(self.stdout_data or self.stderr_data) } def execute_command(command) -> Dict[str, Any]: """Execute a command and return the result.""" executor = CommandExecutor(command) return executor.execute() def get_string(params, key): value = params.get(key) if not isinstance(value, str) or not value.strip(): raise ValueError(f"Invalid or missing field: {key}") return value.strip() @app.route("/api/command", methods=["POST"]) def generic_command(): """Execute any command provided in the request.""" try: params = request.get_json() command = get_string(params, "command") return jsonify(execute_command(command)) except ValueError as e: return jsonify({"success": False, "error": str(e)}), 400 except Exception as e: logger.error(f"Error in command endpoint: {str(e)}") logger.error(traceback.format_exc()) return jsonify({"success": False, "error": str(e)}), 500 @app.route("/api/tools/nmap", methods=["POST"]) def nmap(): """Execute nmap scan with the provided parameters.""" try: params = request.json or {} target = params.get("target", "") scan_type = params.get("scan_type", "-sCV") ports = params.get("ports", "") additional_args = params.get("additional_args", "-T4 -Pn") if not target: logger.warning("Nmap called without target parameter") return jsonify({ "error": "Target parameter is required" }), 400 command = ["nmap"] + shlex.split(scan_type) if ports: command += ["-p", ports] if additional_args: command += shlex.split(additional_args) command.append(target) result = execute_command(command) return jsonify(result) except Exception as e: logger.error(f"Error in nmap endpoint: {str(e)}") logger.error(traceback.format_exc()) return jsonify({ "error": f"Server error: {str(e)}" }), 500 @app.route("/api/tools/gobuster", methods=["POST"]) def gobuster(): """Execute gobuster with the provided parameters.""" try: params = request.json or {} url = params.get("url", "") mode = params.get("mode", "dir") wordlist = params.get("wordlist", "/usr/share/wordlists/dirb/common.txt") additional_args = params.get("additional_args", "") if not url: logger.warning("Gobuster called without URL parameter") return jsonify({ "error": "URL parameter is required" }), 400 # Validate mode if mode not in ["dir", "dns", "fuzz", "vhost"]: logger.warning(f"Invalid gobuster mode: {mode}") return jsonify({ "error": f"Invalid mode: {mode}. Must be one of: dir, dns, fuzz, vhost" }), 400 command = ["gobuster", mode, "-u", url, "-w", wordlist] if additional_args: command += shlex.split(additional_args) result = execute_command(command) return jsonify(result) except Exception as e: logger.error(f"Error in gobuster endpoint: {str(e)}") logger.error(traceback.format_exc()) return jsonify({ "error": f"Server error: {str(e)}" }), 500 # Fix: Added the missing script runner block to execute Flask application if __name__ == "__main__": logger.info(f"Starting Kali API server on port {API_PORT}...") app.run(host="0.0.0.0", port=API_PORT, debug=DEBUG_MODE)