mirror of
https://github.com/splunk/DECEIVE.git
synced 2025-07-02 00:57:26 -04:00
Compare commits
5 Commits
config-fil
...
ssh-execut
Author | SHA1 | Date | |
---|---|---|---|
0197b8b1df | |||
5f27aeeabb | |||
585ee66009 | |||
7be73a7dff | |||
788bd26845 |
@ -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
|
||||
|
@ -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,60 +344,61 @@ 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):
|
||||
# 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"))
|
||||
# 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 our configuration file
|
||||
config = ConfigParser()
|
||||
config.read(args.config)
|
||||
|
||||
# Read the user accounts from the configuration file
|
||||
accounts = get_user_accounts()
|
||||
# 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)
|
||||
# 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 = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
|
||||
logger.addHandler(log_file_handler)
|
||||
|
||||
log_file_handler.setFormatter(JSONFormatter())
|
||||
log_file_handler.setFormatter(JSONFormatter())
|
||||
|
||||
f = ContextFilter()
|
||||
logger.addFilter(f)
|
||||
f = ContextFilter()
|
||||
logger.addFilter(f)
|
||||
|
||||
# Now get access to the LLM
|
||||
# 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"]
|
||||
prompts = get_prompts(args.prompt, args.prompt_file)
|
||||
llm_system_prompt = prompts["system_prompt"]
|
||||
llm_user_prompt = prompts["user_prompt"]
|
||||
|
||||
llm = choose_llm()
|
||||
llm = choose_llm()
|
||||
|
||||
llm_sessions = dict()
|
||||
llm_sessions = dict()
|
||||
|
||||
llm_trimmer = trim_messages(
|
||||
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(
|
||||
llm_prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
"system",
|
||||
@ -377,25 +410,30 @@ llm_prompt = ChatPromptTemplate.from_messages(
|
||||
),
|
||||
MessagesPlaceholder(variable_name="messages"),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
llm_chain = (
|
||||
llm_chain = (
|
||||
RunnablePassthrough.assign(messages=itemgetter("messages") | llm_trimmer)
|
||||
| llm_prompt
|
||||
| llm
|
||||
)
|
||||
)
|
||||
|
||||
with_message_history = RunnableWithMessageHistory(
|
||||
with_message_history = RunnableWithMessageHistory(
|
||||
llm_chain,
|
||||
llm_get_session_history,
|
||||
input_messages_key="messages"
|
||||
)
|
||||
# Thread-local storage for connection details
|
||||
thread_local = threading.local()
|
||||
)
|
||||
# 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()
|
||||
# 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)
|
||||
|
||||
|
Reference in New Issue
Block a user