Merge pull request #6 from DullJZ/main

Support shell switch in meterpreter session

MetasploitMCP.py

@@ -29,11 +29,12 @@ logging.basicConfig(
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger("metasploit_mcp_server")
+session_shell_type: Dict[str, str] = {}
 
 # Metasploit Connection Config (from environment variables)
-MSF_PASSWORD = os.environ.get('MSF_PASSWORD', 'yourpassword')
-MSF_SERVER = os.getenv('MSF_SERVER', '127.0.0.1')
-MSF_PORT_STR = os.getenv('MSF_PORT', '55553')
+MSF_PASSWORD = os.environ.get('MSF_PASSWORD', 'tiaozhanbei')
+MSF_SERVER = os.getenv('MSF_SERVER', '120.24.89.226')
+MSF_PORT_STR = os.getenv('MSF_PORT', '28739')
 MSF_SSL_STR = os.getenv('MSF_SSL', 'false')
 PAYLOAD_SAVE_DIR = os.environ.get('PAYLOAD_SAVE_DIR', str(pathlib.Path.home() / "payloads"))
 
@@ -1039,6 +1040,9 @@ async def send_session_command(
     Uses session.run_with_output for Meterpreter, and a prompt-aware loop for shells.
     The agent is responsible for parsing the raw output.
 
+    In Meterpreter mode, to run a shell command, run `shell` to enter the shell mode first.
+    To exit shell mode and return to Meterpreter, run `exit`.
+
     Args:
         session_id: ID of the target session.
         command: Command string to execute in the session.
@@ -1073,13 +1077,31 @@ async def send_session_command(
         message = "Command execution failed or type unknown."
 
         if session_type == 'meterpreter':
+            if session_shell_type.get(session_id_str) is None:
+                session_shell_type[session_id_str] = 'meterpreter'
+
             logger.debug(f"Using session.run_with_output for Meterpreter session {session_id}")
             try:
                 # Use asyncio.wait_for to handle timeout manually since run_with_output doesn't support timeout parameter
-                output = await asyncio.wait_for(
-                    asyncio.to_thread(lambda: session.run_with_output(command)),
-                    timeout=timeout_seconds
-                )
+                if command == "shell":
+                    if session_shell_type[session_id_str] == 'meterpreter':
+                        output = session.run_with_output(command, end_strs=['created.'])
+                        session_shell_type[session_id_str] = 'shell'
+                        session.read()  # Clear buffer
+                    else:
+                        output = "You are already in shell mode."
+                elif command == "exit":
+                    if session_shell_type[session_id_str] == 'meterpreter':
+                        output = "You are already in meterpreter mode. No need to exit."
+                    else:
+                        session.read()  # Clear buffer
+                        session.detach()
+                        session_shell_type[session_id_str] = 'meterpreter'
+                else:
+                    output = await asyncio.wait_for(
+                        asyncio.to_thread(lambda: session.run_with_output(command)),
+                        timeout=timeout_seconds
+                    )
                 status = "success"
                 message = "Meterpreter command executed successfully."
                 logger.debug(f"Meterpreter command '{command}' completed.")