From 675b7f2f9bfa1a67a143a2d9d273b1388b6b4977 Mon Sep 17 00:00:00 2001
From: DullJZ <79080562+DullJZ@users.noreply.github.com>
Date: Mon, 14 Jul 2025 07:18:16 +0000
Subject: [PATCH 1/2] Support shell switch in meterpreter session

---
 MetasploitMCP.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/MetasploitMCP.py b/MetasploitMCP.py
index a6d971e..0cd770a 100644
--- a/MetasploitMCP.py
+++ b/MetasploitMCP.py
@@ -1076,10 +1076,17 @@ async def send_session_command(
             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":
+                    output = session.run_with_output(command, end_strs=['created.'])
+                    session.read()  # Clear buffer
+                elif command == "exit":
+                    session.read()  # Clear buffer
+                    session.detach()
+                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.")

From 275bbd36c299ac3e6e2fa779846e3407fccc19a3 Mon Sep 17 00:00:00 2001
From: DullJZ <79080562+DullJZ@users.noreply.github.com>
Date: Mon, 14 Jul 2025 08:05:08 +0000
Subject: [PATCH 2/2] Prevent repeated switching of shell modes.

---
 MetasploitMCP.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/MetasploitMCP.py b/MetasploitMCP.py
index 0cd770a..84f553d 100644
--- a/MetasploitMCP.py
+++ b/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,15 +1077,26 @@ 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
                 if command == "shell":
-                    output = session.run_with_output(command, end_strs=['created.'])
-                    session.read()  # Clear buffer
+                    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":
-                    session.read()  # Clear buffer
-                    session.detach()
+                    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)),