5 Commits

Author SHA1 Message Date
0197b8b1df Update config.ini.TEMPLATE to support non-interactive commands
Oops, I forgot to commit this before.
2025-02-04 14:55:02 -05:00
5f27aeeabb Correctly handle both interactive and non-interactive SSH sessions
SSH servers can take user commands from an interactive session as normal, but users can also include commands on the ssh client command line which are executed on the server (e.g., "ssh <hostname> 'uname -a'"). We now execute these non-interactive commands properly as well.

Also added a new "interactive" flag to all user commands (true/false) to show which type of command execution this was.
2025-02-04 12:29:12 -05:00
585ee66009 Don't print ConnectionLost exceptions to the console
These are far too frequent. We still log them, though, we just don't print them.
2025-01-28 10:48:29 -05:00
7be73a7dff Make peername and sockname calls more robust across platforms
For whatever reason, MacOS returns 4 values from conn.get_extra_info('peername') and conn.get_extra_info('sockname'), but Linux systems only return 2.  On the Mac, it's only the first two that we need anyway. Now we retrieve them all, no matter how many there are, and just use the first two so it will work on both platforms.
2025-01-28 10:39:12 -05:00
788bd26845 Now print exceptions to console when SSH connection is lost 2025-01-28 10:21:27 -05:00
2 changed files with 159 additions and 120 deletions

View File

@ -51,10 +51,11 @@ system_prompt = Interpret all inputs as though they were SSH commands and provid
responses do not indicate that the user is on a virtual machine, container,
or emulation at all.
On the first call, be sure to include a realistic MOTD.
If {interactive} is True, make sure to include a realistic MOTD on the
first call, and to end all responses with a realistic shell prompt to display to
the user, including a space at the end.
End all responses with a realistic shell prompt to display to the user,
including a space at the end.
If {interactive} is False, do not include a shell prompt or MOTD.
Include ANSI color codes for the terminal with the output of ls commands
(including any flags), or in any other situation where it is appropriate, but

View File

@ -8,6 +8,7 @@ import threading
import sys
import json
import os
import traceback
from typing import Optional
import logging
import datetime
@ -22,6 +23,7 @@ from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMess
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from asyncssh.misc import ConnectionLost
class JSONFormatter(logging.Formatter):
def format(self, record):
@ -35,6 +37,8 @@ class JSONFormatter(logging.Formatter):
"dst_port": record.dst_port,
"message": record.getMessage()
}
if hasattr(record, 'interactive'):
log_record["interactive"] = record.interactive
# Include any additional fields from the extra dictionary
for key, value in record.__dict__.items():
if key not in log_record and key != 'args' and key != 'msg':
@ -48,8 +52,18 @@ class MySSHServer(asyncssh.SSHServer):
def connection_made(self, conn: asyncssh.SSHServerConnection) -> None:
# Get the source and destination IPs and ports
(src_ip, src_port, _, _) = conn.get_extra_info('peername')
(dst_ip, dst_port, _, _) = conn.get_extra_info('sockname')
peername = conn.get_extra_info('peername')
sockname = conn.get_extra_info('sockname')
if peername is not None:
src_ip, src_port = peername[:2]
else:
src_ip, src_port = '-', '-'
if sockname is not None:
dst_ip, dst_port = sockname[:2]
else:
dst_ip, dst_port = '-', '-'
# Store the connection details in thread-local storage
thread_local.src_ip = src_ip
@ -63,6 +77,8 @@ class MySSHServer(asyncssh.SSHServer):
def connection_lost(self, exc: Optional[Exception]) -> None:
if exc:
logger.error('SSH connection error', extra={"error": str(exc)})
if not isinstance(exc, ConnectionLost):
traceback.print_exception(exc)
else:
logger.info("SSH connection closed")
# Ensure session summary is called on connection loss if attributes are set
@ -136,7 +152,8 @@ representative examples.
llm_response = await session.ainvoke(
{
"messages": [HumanMessage(content=prompt)],
"username": process.get_extra_info('username')
"username": process.get_extra_info('username'),
"interactive": True # Ensure interactive flag is passed
},
config=llm_config
)
@ -164,32 +181,47 @@ async def handle_client(process: asyncssh.SSHServerProcess, server: MySSHServer)
llm_config = {"configurable": {"session_id": task_uuid}}
try:
if process.command:
# Handle non-interactive command execution
command = process.command
logger.info("User input", extra={"details": b64encode(command.encode('utf-8')).decode('utf-8'), "interactive": False})
llm_response = await with_message_history.ainvoke(
{
"messages": [HumanMessage(content=command)],
"username": process.get_extra_info('username'),
"interactive": False
},
config=llm_config
)
process.stdout.write(f"{llm_response.content}")
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": False})
await session_summary(process, llm_config, with_message_history, server)
process.exit(0)
else:
# Handle interactive session
llm_response = await with_message_history.ainvoke(
{
"messages": [HumanMessage(content="ignore this message")],
"username": process.get_extra_info('username')
"username": process.get_extra_info('username'),
"interactive": True
},
config=llm_config
)
process.stdout.write(f"{llm_response.content}")
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8')})
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": True})
# Store process, llm_config, and session in the MySSHServer instance
server._process = process
server._llm_config = llm_config
server._session = with_message_history
try:
async for line in process.stdin:
line = line.rstrip('\n')
logger.info("User input", extra={"details": b64encode(line.encode('utf-8')).decode('utf-8')})
logger.info("User input", extra={"details": b64encode(line.encode('utf-8')).decode('utf-8'), "interactive": True})
# Send the command to the LLM and give the response to the user
llm_response = await with_message_history.ainvoke(
{
"messages": [HumanMessage(content=line)],
"username": process.get_extra_info('username')
"username": process.get_extra_info('username'),
"interactive": True
},
config=llm_config
)
@ -199,7 +231,7 @@ async def handle_client(process: asyncssh.SSHServerProcess, server: MySSHServer)
return
else:
process.stdout.write(f"{llm_response.content}")
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8')})
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": True})
except asyncssh.BreakReceived:
pass
@ -312,6 +344,7 @@ def get_prompts(prompt: Optional[str], prompt_file: Optional[str]) -> dict:
#### MAIN ####
try:
# Parse command line arguments
parser = argparse.ArgumentParser(description='Start the SSH honeypot server.')
parser.add_argument('-c', '--config', type=str, default='config.ini', help='Path to the configuration file')
@ -399,3 +432,8 @@ asyncio.set_event_loop(loop)
loop.run_until_complete(start_server())
loop.run_forever()
except Exception as e:
print(f"Error: {e}", file=sys.stderr)
traceback.print_exc()
sys.exit(1)