CalvinBackup/NeuroSploit
1
0
Fork 0
mirror of https://github.com/CyberSecurityUP/NeuroSploit.git synced 2026-02-12 14:02:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update base_agent.py

Browse Source
This commit is contained in:
Joas A Santos
2026-01-02 11:51:24 -03:00
committed by GitHub
parent 744c1f5113
commit 9aab47c4fc
1 changed files with 32 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
agents/base_agent.py
View File

@@ -36,14 +36,35 @@ class BaseAgent:
logger.warning(f"No user prompt template found for agent {self.agent_name}.")
return user_input # Fallback to raw user input
# Create a dictionary with all the possible placeholders
# Create a dictionary with all the possible placeholders and default values
format_dict = {
"user_input": user_input,
"target_info_json": user_input, # for bug_bounty_hunter
"recon_data_json": json.dumps(additional_context or {}, indent=2), # for bug_bounty_hunter
# For bug_bounty_hunter agent
"target_info_json": user_input,
"recon_data_json": json.dumps(additional_context or {}, indent=2),
# For red_team_agent
"mission_objectives_json": user_input,
"target_environment_json": json.dumps(additional_context or {}, indent=2),
# For pentest agent
"scope_json": user_input,
"initial_info_json": json.dumps(additional_context or {}, indent=2),
# For blue_team_agent
"logs_alerts_json": user_input,
"telemetry_json": json.dumps(additional_context or {}, indent=2),
# For exploit_expert agent
"vulnerability_details_json": user_input,
"target_info_json": json.dumps(additional_context or {}, indent=2),
# For cwe_expert agent
"code_vulnerability_json": user_input,
# For malware_analysis agent
"malware_sample_json": user_input,
# For replay_attack agent
"traffic_logs_json": user_input,
# Generic additional context
"additional_context_json": json.dumps(additional_context or {}, indent=2)
}
# Override with actual additional_context values if provided
if additional_context:
for key, value in additional_context.items():
if isinstance(value, (dict, list)):
@@ -51,9 +72,14 @@ class BaseAgent:
else:
format_dict[key] = value
# Use a safe way to format, ignoring missing keys
formatted_prompt = user_prompt_template.format_map({k: v for k, v in format_dict.items() if f"{{{k}}}" in user_prompt_template})
# Use a safe way to format, only including keys that exist in the template
try:
formatted_prompt = user_prompt_template.format_map(format_dict)
except KeyError as e:
logger.error(f"Missing key in format_dict: {e}")
# Fallback to user input if formatting fails
return user_input
return formatted_prompt
def execute(self, user_input: str, campaign_data: Dict = None) -> Dict: