Moved SSH honeypot to subdirectory

SSH/ssh_server.py

@@ -0,0 +1,265 @@
+#!/usr/bin/env python3
+
+from configparser import ConfigParser
+import asyncio
+import asyncssh
+import threading
+import sys
+import json
+from typing import Optional
+import logging
+import datetime
+import uuid
+from base64 import b64encode
+from operator import itemgetter
+from langchain_openai import ChatOpenAI
+from langchain_aws import ChatBedrock, ChatBedrockConverse
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import HumanMessage, SystemMessage, trim_messages
+from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnablePassthrough
+
+async def handle_client(process: asyncssh.SSHServerProcess) -> None:
+# This is the main loop for handling SSH client connections. 
+# Any user interaction should be done here.
+
+    # Give each session a unique name
+    task_uuid = f"session-{uuid.uuid4()}"
+    current_task = asyncio.current_task()
+    current_task.set_name(task_uuid)
+
+    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(f"OUTPUT: {b64encode(llm_response.content.encode('ascii')).decode('ascii')}")
+
+    try:
+        async for line in process.stdin:
+            line = line.rstrip('\n')
+            logger.info(f"INPUT: {line}")
+
+            # 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')
+                },
+                    config=llm_config
+            )
+            if llm_response.content == "XXX-END-OF-SESSION-XXX":
+                process.exit(0)
+            else:
+                process.stdout.write(f"{llm_response.content}")
+                logger.info(f"OUTPUT: {b64encode(llm_response.content.encode('ascii')).decode('ascii')}")
+
+    except asyncssh.BreakReceived:
+        pass
+
+    # Just in case we ever get here, which we probably shouldn't
+    process.exit(0)
+
+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')
+
+        # Store the connection details in thread-local storage
+        thread_local.src_ip = src_ip
+        thread_local.src_port = src_port
+        thread_local.dst_ip = dst_ip
+        thread_local.dst_port = dst_port
+
+        # Log the connection details
+        logger.info(f"SSH connection received from {src_ip}/{src_port} to {dst_ip}/{dst_port}.")
+
+    def connection_lost(self, exc: Optional[Exception]) -> None:
+        if exc:
+            logger.error('SSH connection error: ' + str(exc))
+
+        else:
+            logger.info("SSH connection closed.")
+
+    def begin_auth(self, username: str) -> bool:
+        if accounts.get(username) != '':
+            logger.info(f"AUTH: User {username} attempting to authenticate.")
+            return True
+        else:
+            logger.info(f"AUTH: SUCCESS for user {username} with password ''.")
+            return False
+
+    def password_auth_supported(self) -> bool:
+        return True
+    def host_based_auth_supported(self) -> bool:
+        return False
+    def public_key_auth_supported(self) -> bool:
+        return False
+    def kbdinit_auth_supported(self) -> bool:
+        return False
+
+    def validate_password(self, username: str, password: str) -> bool:
+        pw = accounts.get(username, '*')
+        
+        if ((pw != '*') and (password == pw)):
+            logger.info(f"AUTH: SUCCESS for user {username} with password '{password}'.")
+            return True
+        else:
+            logger.info(f"AUTH: FAILED for user {username} with password '{password}'.")
+            return False
+
+async def start_server() -> None:
+    await asyncssh.listen(
+        port=config['ssh'].getint("port", 8022),
+        reuse_address=True,
+        reuse_port=True,
+        server_factory=MySSHServer,
+        server_host_keys=config['ssh'].get("host_priv_key", "ssh_host_key"),
+        process_factory=handle_client,
+        server_version=config['ssh'].get("server_version_string", "SSH-2.0-OpenSSH_8.2p1 Ubuntu-4ubuntu0.3")
+    )
+
+class ContextFilter(logging.Filter):
+    """
+    This filter is used to add the current asyncio task name to the log record,
+    so you can group events in the same session together.
+    """
+
+    def filter(self, record):
+
+        task = asyncio.current_task()
+        if task:
+            task_name = task.get_name()
+        else:
+            task_name = "-"
+
+        record.src_ip = thread_local.__dict__.get('src_ip', '-')
+        record.src_port = thread_local.__dict__.get('src_port', '-')   
+        record.dst_ip = thread_local.__dict__.get('dst_ip', '-')
+        record.dst_port = thread_local.__dict__.get('dst_port', '-')
+
+        record.task_name = task_name
+
+        return True
+
+def llm_get_session_history(session_id: str) -> BaseChatMessageHistory:
+    if session_id not in llm_sessions:
+        llm_sessions[session_id] = InMemoryChatMessageHistory()
+    return llm_sessions[session_id]
+
+def get_user_accounts() -> dict:
+    if (not 'user_accounts' in config) or (len(config.items('user_accounts')) == 0):
+        raise ValueError("No user accounts found in configuration file.")
+    
+    accounts = dict()
+
+    for k, v in config.items('user_accounts'):
+        accounts[k] = v
+
+    return accounts
+
+def choose_llm():
+    llm_provider_name = config['llm'].get("llm_provider", "openai")
+    llm_provider_name = llm_provider_name.lower()
+    model_name = config['llm'].get("model_name", "gpt-3.5-turbo")
+
+    if llm_provider_name == 'openai':
+        llm_model = ChatOpenAI(
+            model=model_name
+        )
+    elif llm_provider_name == 'aws':
+        llm_model = ChatBedrockConverse(
+            model=model_name,
+            region_name=config['llm'].get("aws_region", "us-east-1"),
+            credentials_profile_name=config['llm'].get("aws_credentials_profile", "default")        )
+    elif llm_provider_name == 'gemini':
+        llm_model = ChatGoogleGenerativeAI(
+            model=model_name,
+        )
+    else:
+        raise ValueError(f"Invalid LLM provider {llm_provider_name}.")
+
+    return llm_model
+
+#### MAIN ####
+
+# Always use UTC for logging
+logging.Formatter.formatTime = (lambda self, record, datefmt=None: datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).astimezone().isoformat(sep="T",timespec="milliseconds"))
+
+# Read our configuration file
+config = ConfigParser()
+config.read("config.ini")
+
+# 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(logging.Formatter("%(asctime)s %(levelname)s %(task_name)s SSH %(src_ip)s:%(src_port)s -> %(dst_ip)s:%(dst_port)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S. %Z"))
+
+f = ContextFilter()
+logger.addFilter(f)
+
+# Now get access to the LLM
+
+prompt_file = config['llm'].get("system_prompt_file", "prompt.txt")
+with open(prompt_file, "r") as f:
+    llm_system_prompt = f.read()
+
+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
+        ),
+        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()
+