splunk/DECEIVE
1
0
Fork 0
mirror of https://github.com/splunk/DECEIVE.git synced 2025-07-01 16:47:28 -04:00
Code Issues Packages Projects Releases Wiki Activity

New command line options for prompts and config files.

Browse Source 
* --prompt-file to specify a file from which to read the prompt.
* --prompt to specify a prompt string on the command line
* --config to specify an alternate config file
This commit is contained in:
David J. Bianco
2025-01-27 13:20:41 -05:00
parent 545d50f294
commit cea5dc28a2
1 changed files with 33 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
SSH/ssh_server.py
View File

@ -1,11 +1,13 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
from configparser import ConfigParser from configparser import ConfigParser
import argparse
import asyncio import asyncio
import asyncssh import asyncssh
import threading import threading
import sys import sys
import json import json
import os
from typing import Optional from typing import Optional
import logging import logging
import datetime import datetime
@ -285,10 +287,24 @@ def choose_llm():
return llm_model return llm_model
def get_prompts() -> dict: def get_prompts(prompt: Optional[str], prompt_file: Optional[str]) -> dict:
system_prompt = config['llm']['system_prompt'] system_prompt = config['llm']['system_prompt']
with open("prompt.txt", "r") as f: if prompt is not None:
user_prompt = f.read() if not prompt.strip():
print("Error: The prompt text cannot be empty.", file=sys.stderr)
sys.exit(1)
user_prompt = prompt
elif prompt_file:
if not os.path.exists(prompt_file):
print(f"Error: The specified prompt file '{prompt_file}' does not exist.", file=sys.stderr)
sys.exit(1)
with open(prompt_file, "r") as f:
user_prompt = f.read()
elif os.path.exists("prompt.txt"):
with open("prompt.txt", "r") as f:
user_prompt = f.read()
else:
raise ValueError("Either prompt or prompt_file must be provided.")
return { return {
"system_prompt": system_prompt, "system_prompt": system_prompt,
"user_prompt": user_prompt "user_prompt": user_prompt
@ -296,12 +312,24 @@ def get_prompts() -> dict:
#### MAIN #### #### 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()
# 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 # 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")) 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 # Read our configuration file
config = ConfigParser() config = ConfigParser()
config.read("config.ini") config.read(args.config)
# Read the user accounts from the configuration file # Read the user accounts from the configuration file
accounts = get_user_accounts() accounts = get_user_accounts()
@ -320,7 +348,7 @@ logger.addFilter(f)
# Now get access to the LLM # Now get access to the LLM
prompts = get_prompts() prompts = get_prompts(args.prompt, args.prompt_file)
llm_system_prompt = prompts["system_prompt"] llm_system_prompt = prompts["system_prompt"]
llm_user_prompt = prompts["user_prompt"] llm_user_prompt = prompts["user_prompt"]