diff --git a/bbot/scanner/preset/args.py b/bbot/scanner/preset/args.py index 11df99022..b39ce0896 100644 --- a/bbot/scanner/preset/args.py +++ b/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() diff --git a/bbot/test/test_step_1/test_cli.py b/bbot/test/test_step_1/test_cli.py index 4b1f13143..1b1cfd035 100644 --- a/bbot/test/test_step_1/test_cli.py +++ b/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)