adding custom headers CLI argument / tests

bbot/scanner/preset/args.py

@@ -145,6 +145,7 @@ def preset_from_args(self):
         args_preset.scan_name = self.parsed.name
         args_preset.output_dir = self.parsed.output_dir
         args_preset.force_start = self.parsed.force
+        args_preset.core.merge_custom({"http_headers": self.parsed.custom_headers})
 
         # CLI config options (dot-syntax)
         for config_arg in self.parsed.config:
@@ -299,6 +300,13 @@ def create_parser(self, *args, **kwargs):
 
         misc = p.add_argument_group(title="Misc")
         misc.add_argument("--version", action="store_true", help="show BBOT version and exit")
+        misc.add_argument(
+            "-H",
+            "--custom-headers",
+            nargs="+",
+            default=[],
+            help="List of custom headers as key value pairs (header=value).",
+        )
         return p
 
     def sanitize_args(self):
@@ -323,6 +331,22 @@ def sanitize_args(self):
         self.parsed.require_flags = chain_lists(self.parsed.require_flags)
         self.parsed.event_types = [t.upper() for t in chain_lists(self.parsed.event_types)]
 
+        # Custom Header Parsing / Validation
+        custom_headers_dict = {}
+        custom_header_example = "Example: --custom-headers foo=bar foo2=bar2"
+
+        for i in self.parsed.custom_headers:
+            parts = i.split("=", 1)
+            if len(parts) != 2:
+                raise ValidationError(f"Custom headers not formatted correctly (missing '='). {custom_header_example}")
+            k, v = parts
+            if not k or not v:
+                raise ValidationError(
+                    f"Custom headers not formatted correctly (missing header name or value). {custom_header_example}"
+                )
+            custom_headers_dict[k] = v
+        self.parsed.custom_headers = custom_headers_dict
+
     def validate(self):
         # validate config options
         sentinel = object()

bbot/test/test_step_1/test_cli.py

@@ -392,6 +392,43 @@ async def test_cli_args(monkeypatch, caplog, capsys, clean_default_config):
     assert success == True, "--install-all-deps failed for at least one module"
 
 
+@pytest.mark.asyncio
+async def test_cli_customheaders(monkeypatch, caplog, capsys):
+    monkeypatch.setattr(sys, "exit", lambda *args, **kwargs: True)
+    monkeypatch.setattr(os, "_exit", lambda *args, **kwargs: True)
+    import yaml
+
+    # test custom headers
+    monkeypatch.setattr(
+        "sys.argv", ["bbot", "--custom-headers", "foo=bar", "foo2=bar2", "foo3=bar=3", "--current-preset"]
+    )
+    success = await cli._main()
+    assert success == None, "setting custom headers on command line failed"
+    captured = capsys.readouterr()
+    stdout_preset = yaml.safe_load(captured.out)
+    assert stdout_preset["config"]["http_headers"] == {"foo": "bar", "foo2": "bar2", "foo3": "bar=3"}
+
+    # test custom headers invalid (no "=")
+    monkeypatch.setattr("sys.argv", ["bbot", "--custom-headers", "justastring", "--current-preset"])
+    result = await cli._main()
+    assert result == None
+    assert "Custom headers not formatted correctly (missing '=')" in caplog.text
+    caplog.clear()
+
+    # test custom headers invalid (missing key)
+    monkeypatch.setattr("sys.argv", ["bbot", "--custom-headers", "=nokey", "--current-preset"])
+    result = await cli._main()
+    assert result == None
+    assert "Custom headers not formatted correctly (missing header name or value)" in caplog.text
+    caplog.clear()
+
+    # test custom headers invalid (missing value)
+    monkeypatch.setattr("sys.argv", ["bbot", "--custom-headers", "missingvalue=", "--current-preset"])
+    result = await cli._main()
+    assert result == None
+    assert "Custom headers not formatted correctly (missing header name or value)" in caplog.text
+
+
 def test_cli_config_validation(monkeypatch, caplog):
     monkeypatch.setattr(sys, "exit", lambda *args, **kwargs: True)
     monkeypatch.setattr(os, "_exit", lambda *args, **kwargs: True)