|
|
|
@ -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
|
|
|
|
@ -89,7 +105,7 @@ class MySSHServer(asyncssh.SSHServer):
|
|
|
|
|
def validate_password(self, username: str, password: str) -> bool:
|
|
|
|
|
pw = accounts.get(username, '*')
|
|
|
|
|
|
|
|
|
|
if ((pw != '*') and (password == pw)):
|
|
|
|
|
if pw == '*' or (pw != '*' and password == pw):
|
|
|
|
|
logger.info("Authentication success", extra={"username": username, "password": password})
|
|
|
|
|
return True
|
|
|
|
|
else:
|
|
|
|
@ -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,42 +181,57 @@ async def handle_client(process: asyncssh.SSHServerProcess, server: MySSHServer)
|
|
|
|
|
|
|
|
|
|
llm_config = {"configurable": {"session_id": task_uuid}}
|
|
|
|
|
|
|
|
|
|
llm_response = await with_message_history.ainvoke(
|
|
|
|
|
{
|
|
|
|
|
"messages": [HumanMessage(content="ignore this message")],
|
|
|
|
|
"username": process.get_extra_info('username')
|
|
|
|
|
},
|
|
|
|
|
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')})
|
|
|
|
|
|
|
|
|
|
# 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')})
|
|
|
|
|
|
|
|
|
|
# Send the command to the LLM and give the response to the user
|
|
|
|
|
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=line)],
|
|
|
|
|
"username": process.get_extra_info('username')
|
|
|
|
|
"messages": [HumanMessage(content=command)],
|
|
|
|
|
"username": process.get_extra_info('username'),
|
|
|
|
|
"interactive": False
|
|
|
|
|
},
|
|
|
|
|
config=llm_config
|
|
|
|
|
)
|
|
|
|
|
if llm_response.content == "XXX-END-OF-SESSION-XXX":
|
|
|
|
|
await session_summary(process, llm_config, with_message_history, server)
|
|
|
|
|
process.exit(0)
|
|
|
|
|
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')})
|
|
|
|
|
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'),
|
|
|
|
|
"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'), "interactive": True})
|
|
|
|
|
|
|
|
|
|
async for line in process.stdin:
|
|
|
|
|
line = line.rstrip('\n')
|
|
|
|
|
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'),
|
|
|
|
|
"interactive": True
|
|
|
|
|
},
|
|
|
|
|
config=llm_config
|
|
|
|
|
)
|
|
|
|
|
if llm_response.content == "XXX-END-OF-SESSION-XXX":
|
|
|
|
|
await session_summary(process, llm_config, with_message_history, server)
|
|
|
|
|
process.exit(0)
|
|
|
|
|
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'), "interactive": True})
|
|
|
|
|
|
|
|
|
|
except asyncssh.BreakReceived:
|
|
|
|
|
pass
|
|
|
|
@ -312,90 +344,96 @@ def get_prompts(prompt: Optional[str], prompt_file: Optional[str]) -> dict:
|
|
|
|
|
|
|
|
|
|
#### MAIN ####
|
|
|
|
|
|
|
|
|
|
# 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')
|
|
|
|
|
parser.add_argument('-p', '--prompt', type=str, help='The entire text of the prompt')
|
|
|
|
|
parser.add_argument('-f', '--prompt-file', type=str, default='prompt.txt', help='Path to the prompt file')
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
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')
|
|
|
|
|
parser.add_argument('-p', '--prompt', type=str, help='The entire text of the prompt')
|
|
|
|
|
parser.add_argument('-f', '--prompt-file', type=str, default='prompt.txt', help='Path to the prompt file')
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
|
|
|
|
# Check if the config file exists
|
|
|
|
|
if not os.path.exists(args.config):
|
|
|
|
|
print(f"Error: The specified config file '{args.config}' does not exist.", file=sys.stderr)
|
|
|
|
|
# Check if the config file exists
|
|
|
|
|
if not os.path.exists(args.config):
|
|
|
|
|
print(f"Error: The specified config file '{args.config}' does not exist.", file=sys.stderr)
|
|
|
|
|
sys.exit(1)
|
|
|
|
|
|
|
|
|
|
# Always use UTC for logging
|
|
|
|
|
logging.Formatter.formatTime = (lambda self, record, datefmt=None: datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).isoformat(sep="T",timespec="milliseconds"))
|
|
|
|
|
|
|
|
|
|
# Read our configuration file
|
|
|
|
|
config = ConfigParser()
|
|
|
|
|
config.read(args.config)
|
|
|
|
|
|
|
|
|
|
# Read the user accounts from the configuration file
|
|
|
|
|
accounts = get_user_accounts()
|
|
|
|
|
|
|
|
|
|
# Set up the honeypot logger
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
|
|
|
|
log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
|
|
|
|
|
logger.addHandler(log_file_handler)
|
|
|
|
|
|
|
|
|
|
log_file_handler.setFormatter(JSONFormatter())
|
|
|
|
|
|
|
|
|
|
f = ContextFilter()
|
|
|
|
|
logger.addFilter(f)
|
|
|
|
|
|
|
|
|
|
# Now get access to the LLM
|
|
|
|
|
|
|
|
|
|
prompts = get_prompts(args.prompt, args.prompt_file)
|
|
|
|
|
llm_system_prompt = prompts["system_prompt"]
|
|
|
|
|
llm_user_prompt = prompts["user_prompt"]
|
|
|
|
|
|
|
|
|
|
llm = choose_llm()
|
|
|
|
|
|
|
|
|
|
llm_sessions = dict()
|
|
|
|
|
|
|
|
|
|
llm_trimmer = trim_messages(
|
|
|
|
|
max_tokens=config['llm'].getint("trimmer_max_tokens", 64000),
|
|
|
|
|
strategy="last",
|
|
|
|
|
token_counter=llm,
|
|
|
|
|
include_system=True,
|
|
|
|
|
allow_partial=False,
|
|
|
|
|
start_on="human",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
llm_prompt = ChatPromptTemplate.from_messages(
|
|
|
|
|
[
|
|
|
|
|
(
|
|
|
|
|
"system",
|
|
|
|
|
llm_system_prompt
|
|
|
|
|
),
|
|
|
|
|
(
|
|
|
|
|
"system",
|
|
|
|
|
llm_user_prompt
|
|
|
|
|
),
|
|
|
|
|
MessagesPlaceholder(variable_name="messages"),
|
|
|
|
|
]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
llm_chain = (
|
|
|
|
|
RunnablePassthrough.assign(messages=itemgetter("messages") | llm_trimmer)
|
|
|
|
|
| llm_prompt
|
|
|
|
|
| llm
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
with_message_history = RunnableWithMessageHistory(
|
|
|
|
|
llm_chain,
|
|
|
|
|
llm_get_session_history,
|
|
|
|
|
input_messages_key="messages"
|
|
|
|
|
)
|
|
|
|
|
# Thread-local storage for connection details
|
|
|
|
|
thread_local = threading.local()
|
|
|
|
|
|
|
|
|
|
# Kick off the server!
|
|
|
|
|
loop = asyncio.new_event_loop()
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
|
|
# Always use UTC for logging
|
|
|
|
|
logging.Formatter.formatTime = (lambda self, record, datefmt=None: datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).isoformat(sep="T",timespec="milliseconds"))
|
|
|
|
|
|
|
|
|
|
# Read our configuration file
|
|
|
|
|
config = ConfigParser()
|
|
|
|
|
config.read(args.config)
|
|
|
|
|
|
|
|
|
|
# Read the user accounts from the configuration file
|
|
|
|
|
accounts = get_user_accounts()
|
|
|
|
|
|
|
|
|
|
# Set up the honeypot logger
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
|
|
|
|
log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
|
|
|
|
|
logger.addHandler(log_file_handler)
|
|
|
|
|
|
|
|
|
|
log_file_handler.setFormatter(JSONFormatter())
|
|
|
|
|
|
|
|
|
|
f = ContextFilter()
|
|
|
|
|
logger.addFilter(f)
|
|
|
|
|
|
|
|
|
|
# Now get access to the LLM
|
|
|
|
|
|
|
|
|
|
prompts = get_prompts(args.prompt, args.prompt_file)
|
|
|
|
|
llm_system_prompt = prompts["system_prompt"]
|
|
|
|
|
llm_user_prompt = prompts["user_prompt"]
|
|
|
|
|
|
|
|
|
|
llm = choose_llm()
|
|
|
|
|
|
|
|
|
|
llm_sessions = dict()
|
|
|
|
|
|
|
|
|
|
llm_trimmer = trim_messages(
|
|
|
|
|
max_tokens=config['llm'].getint("trimmer_max_tokens", 64000),
|
|
|
|
|
strategy="last",
|
|
|
|
|
token_counter=llm,
|
|
|
|
|
include_system=True,
|
|
|
|
|
allow_partial=False,
|
|
|
|
|
start_on="human",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
llm_prompt = ChatPromptTemplate.from_messages(
|
|
|
|
|
[
|
|
|
|
|
(
|
|
|
|
|
"system",
|
|
|
|
|
llm_system_prompt
|
|
|
|
|
),
|
|
|
|
|
(
|
|
|
|
|
"system",
|
|
|
|
|
llm_user_prompt
|
|
|
|
|
),
|
|
|
|
|
MessagesPlaceholder(variable_name="messages"),
|
|
|
|
|
]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
llm_chain = (
|
|
|
|
|
RunnablePassthrough.assign(messages=itemgetter("messages") | llm_trimmer)
|
|
|
|
|
| llm_prompt
|
|
|
|
|
| llm
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
with_message_history = RunnableWithMessageHistory(
|
|
|
|
|
llm_chain,
|
|
|
|
|
llm_get_session_history,
|
|
|
|
|
input_messages_key="messages"
|
|
|
|
|
)
|
|
|
|
|
# Thread-local storage for connection details
|
|
|
|
|
thread_local = threading.local()
|
|
|
|
|
|
|
|
|
|
# Kick off the server!
|
|
|
|
|
loop = asyncio.new_event_loop()
|
|
|
|
|
asyncio.set_event_loop(loop)
|
|
|
|
|
loop.run_until_complete(start_server())
|
|
|
|
|
loop.run_forever()
|
|
|
|
|
|
|
|
|
|