Merge pull request #2190 from blacklanternsecurity/custom-ua

Add --user-agent and -user-agent-suffix CLI options

bbot/scanner/preset/args.py

@@ -168,6 +168,21 @@ def preset_from_args(self):
                 {"modules": {"excavate": {"custom_yara_rules": self.parsed.custom_yara_rules}}}
             )
 
+        # Check if both user_agent and user_agent_suffix are set. If so combine them and merge into the config
+        if self.parsed.user_agent and self.parsed.user_agent_suffix:
+            modified_user_agent = f"{self.parsed.user_agent} {self.parsed.user_agent_suffix}"
+            args_preset.core.merge_custom({"web": {"user_agent": modified_user_agent}})
+
+        # If only user_agent_suffix is set, retrieve the existing user_agent from the merged config and append the suffix
+        elif self.parsed.user_agent_suffix:
+            existing_user_agent = args_preset.core.config.get("web", {}).get("user_agent", "")
+            modified_user_agent = f"{existing_user_agent} {self.parsed.user_agent_suffix}"
+            args_preset.core.merge_custom({"web": {"user_agent": modified_user_agent}})
+
+        # If only user_agent is set, merge it directly
+        elif self.parsed.user_agent:
+            args_preset.core.merge_custom({"web": {"user_agent": self.parsed.user_agent}})
+
         # CLI config options (dot-syntax)
         for config_arg in self.parsed.config:
             try:
@@ -348,6 +363,15 @@ def create_parser(self, *args, **kwargs):
             help="List of custom headers as key value pairs (header=value).",
         )
         misc.add_argument("--custom-yara-rules", "-cy", help="Add custom yara rules to excavate")
+
+        misc.add_argument("--user-agent", "-ua", help="Set the user-agent for all HTTP requests")
+        misc.add_argument(
+            "--user-agent-suffix",
+            "-uas",
+            help=argparse.SUPPRESS,
+            metavar="SUFFIX",
+            default=None
+        )
         return p
 
     def sanitize_args(self):