From 492ce5ea41ca3d3cc38e9e9112f824ca67703bed Mon Sep 17 00:00:00 2001 From: MayersScott Date: Thu, 7 May 2026 14:19:41 +0300 Subject: [PATCH] feat: support ad-hoc URL probing via --url (can repeat) --- rkn_checker/cli.py | 121 +++++++++++++++++++++++++++++++++++++--- tests/test_cli.py | 108 ++++++++++++++++++++++++++++++++++- tests/test_streaming.py | 4 +- 3 files changed, 223 insertions(+), 10 deletions(-) diff --git a/rkn_checker/cli.py b/rkn_checker/cli.py index f03ef56..fdc75a3 100644 --- a/rkn_checker/cli.py +++ b/rkn_checker/cli.py @@ -5,6 +5,7 @@ import json import logging import sys from concurrent.futures import ThreadPoolExecutor, as_completed +from urllib.parse import urlparse from .core import check_url, get_self_info from .lists import ListLoadError, load_targets @@ -33,6 +34,12 @@ def _build_parser() -> argparse.ArgumentParser: p.add_argument("--black-file", dest="black_file", metavar="PATH", help="load blacklist targets from a .txt or .json file " "(replaces the built-in blacklist)") + p.add_argument("--url", dest="urls", action="append", default=[], + metavar="URL", + help="probe a single URL or hostname; repeat to probe " + "multiple. Skips the built-in lists entirely; runs as " + "an ad-hoc check with no whitelist/blacklist semantics. " + "Conflicts with --white-file/--black-file.") p.add_argument("--timeout", type=float, default=5.0, help="per-probe timeout in seconds (default: 5.0)") p.add_argument("--workers", type=int, default=10, @@ -41,6 +48,12 @@ def _build_parser() -> argparse.ArgumentParser: help="increase log verbosity (-v info, -vv debug)") p.add_argument("--no-self-info", dest="no_self_info", action="store_true", help="skip the external IP self-info lookup") + p.add_argument("--identify", dest="identify", action="store_true", + help="send a self-identifying User-Agent ('rkn-block-checker/') " + "instead of a generic Chrome UA. Use this when probing " + "infrastructure you control or have permission to probe; " + "the default blends in with normal traffic to avoid " + "leaving a unique fingerprint in network logs.") return p @@ -70,6 +83,32 @@ def _resolve_lists( return white, black +def _ad_hoc_targets(raw_urls: list[str]) -> dict[str, str]: + """Turn `--url X --url Y` into a {name: url} mapping for the probe pipeline. + + Unlike list-file loaders, this is forgiving on input: bare hostnames get + https:// prepended, and names are derived from the hostname so the user + doesn't have to invent them. Duplicate hostnames are de-duplicated by + suffixing -2, -3, ... so all explicit `--url` arguments survive. + """ + out: dict[str, str] = {} + for raw in raw_urls: + url = raw.strip() + if not url: + continue + if "://" not in url: + url = "https://" + url + host = urlparse(url).hostname or url + base_name = host.replace(".", "-") + name = base_name + n = 2 + while name in out: + name = f"{base_name}-{n}" + n += 1 + out[name] = url + return out + + def _run_streaming( run_white: bool, run_black: bool, @@ -77,17 +116,18 @@ def _run_streaming( black_urls: dict[str, str], workers: int, timeout: float, + identify: bool = False, ) -> tuple[list[CheckResult], list[CheckResult]]: white_results: list[CheckResult] = [] black_results: list[CheckResult] = [] with ThreadPoolExecutor(max_workers=workers) as pool: white_futs = { - pool.submit(check_url, name, url, timeout): name + pool.submit(check_url, name, url, timeout, identify): name for name, url in (white_urls.items() if run_white else []) } black_futs = { - pool.submit(check_url, name, url, timeout): name + pool.submit(check_url, name, url, timeout, identify): name for name, url in (black_urls.items() if run_black else []) } @@ -110,6 +150,28 @@ def _run_streaming( return white_results, black_results +def _run_ad_hoc( + targets: dict[str, str], + workers: int, + timeout: float, + identify: bool, +) -> list[CheckResult]: + """Probe an ad-hoc list of URLs in a single section, no white/black split.""" + results: list[CheckResult] = [] + print_section(f"Ad-hoc URLs ({len(targets)})") + with ThreadPoolExecutor(max_workers=workers) as pool: + futs = { + pool.submit(check_url, name, url, timeout, identify): name + for name, url in targets.items() + } + for fut in as_completed(futs): + r = fut.result() + results.append(r) + print_result(r) + sys.stdout.flush() + return results + + def main(argv: list[str] | None = None) -> int: parser = _build_parser() args = parser.parse_args(argv) @@ -121,6 +183,46 @@ def main(argv: list[str] | None = None) -> int: parser.error("--timeout must be positive") if args.workers <= 0: parser.error("--workers must be positive") + if args.urls and (args.white_file or args.black_file + or args.white_only or args.black_only): + parser.error( + "--url cannot be combined with --white/--black/--white-file/--black-file; " + "ad-hoc mode runs the listed URLs and nothing else" + ) + + # Ad-hoc single-URL mode: skip the whitelist/blacklist plumbing entirely. + # No summary verdict is produced — there's no control group to compare + # against, and `--url` is for "did this one site come through?" usage, + # not for diagnosing whether the user is in a blocked zone. + if args.urls: + ad_hoc = _ad_hoc_targets(args.urls) + if not ad_hoc: + parser.error("no usable --url targets after parsing") + + if args.as_json: + from .core import check_urls_parallel + results = check_urls_parallel( + ad_hoc, args.workers, args.timeout, identify=args.identify, + ) + self_info = ( + get_self_info(timeout=args.timeout) + if not args.no_self_info else None + ) + payload = { + "self_info": self_info, + "ad_hoc": [r.to_dict() for r in results], + } + json.dump(payload, sys.stdout, ensure_ascii=False, indent=2) + sys.stdout.write("\n") + return 0 + + if not args.no_self_info: + print_header(get_self_info(timeout=args.timeout)) + else: + print_header({}) + sys.stdout.flush() + _run_ad_hoc(ad_hoc, args.workers, args.timeout, args.identify) + return 0 try: white_urls, black_urls = _resolve_lists(args.white_file, args.black_file) @@ -134,11 +236,15 @@ def main(argv: list[str] | None = None) -> int: if args.as_json: from .core import check_urls_parallel white_results = ( - check_urls_parallel(white_urls, args.workers, args.timeout) + check_urls_parallel( + white_urls, args.workers, args.timeout, identify=args.identify, + ) if run_white else [] ) black_results = ( - check_urls_parallel(black_urls, args.workers, args.timeout) + check_urls_parallel( + black_urls, args.workers, args.timeout, identify=args.identify, + ) if run_black else [] ) self_info = get_self_info(timeout=args.timeout) if not args.no_self_info else None @@ -150,7 +256,7 @@ def main(argv: list[str] | None = None) -> int: json.dump(payload, sys.stdout, ensure_ascii=False, indent=2) sys.stdout.write("\n") return 0 - + if not args.no_self_info: print_header(get_self_info(timeout=args.timeout)) else: @@ -158,7 +264,8 @@ def main(argv: list[str] | None = None) -> int: sys.stdout.flush() white_results, black_results = _run_streaming( - run_white, run_black, white_urls, black_urls, args.workers, args.timeout, + run_white, run_black, white_urls, black_urls, + args.workers, args.timeout, args.identify, ) if run_white and run_black: @@ -168,4 +275,4 @@ def main(argv: list[str] | None = None) -> int: if __name__ == "__main__": - sys.exit(main()) \ No newline at end of file + sys.exit(main()) diff --git a/tests/test_cli.py b/tests/test_cli.py index d19da38..7591c44 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from rkn_checker.cli import main +from rkn_checker.cli import _ad_hoc_targets, main class TestMutuallyExclusiveFlags: @@ -44,3 +44,109 @@ class TestStreamingNoSelfInfo: def test_no_self_info_passes_empty_dict_to_header(self, mock_stream, mock_header): main(["--no-self-info"]) mock_header.assert_called_with({}) + + +class TestAdHocTargets: + """Unit tests for the --url -> {name: url} mapping. The mapping must + accept casual user input (bare hostnames, mixed case) and never silently + drop a target the user explicitly asked for.""" + + def test_single_url_with_scheme(self): + out = _ad_hoc_targets(["https://example.com"]) + assert out == {"example-com": "https://example.com"} + + def test_bare_hostname_gets_https_prepended(self): + out = _ad_hoc_targets(["example.com"]) + assert out == {"example-com": "https://example.com"} + + def test_multiple_urls_all_kept(self): + out = _ad_hoc_targets([ + "https://example.com", + "https://other.org", + ]) + assert len(out) == 2 + assert "https://example.com" in out.values() + assert "https://other.org" in out.values() + + def test_duplicate_hosts_disambiguated_not_lost(self): + # Two different URLs that happen to share a hostname (e.g., http vs + # https for the same host) should both end up in the output, not + # collide and silently drop one. + out = _ad_hoc_targets([ + "https://example.com", + "http://example.com", + ]) + assert len(out) == 2 + + def test_empty_string_skipped(self): + out = _ad_hoc_targets(["", "https://x.com", " "]) + assert len(out) == 1 + assert "https://x.com" in out.values() + + +class TestAdHocCli: + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + @patch("rkn_checker.cli._run_ad_hoc", return_value=[]) + def test_url_flag_runs_ad_hoc_path(self, mock_ad_hoc, mock_self): + main(["--url", "https://example.com"]) + mock_ad_hoc.assert_called_once() + + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + @patch("rkn_checker.cli._run_streaming", return_value=([], [])) + def test_url_flag_skips_streaming_path(self, mock_stream, mock_self): + # When --url is used, the regular whitelist/blacklist pipeline must + # not run at all — otherwise the user pays the cost of probing + # ~30 built-in sites for a single ad-hoc check. + main(["--url", "https://example.com"]) + mock_stream.assert_not_called() + + def test_url_with_white_file_rejected(self): + with pytest.raises(SystemExit) as ei: + main(["--url", "https://example.com", + "--white-file", "/tmp/nonexistent.txt"]) + assert ei.value.code == 2 + + def test_url_with_white_only_flag_rejected(self): + with pytest.raises(SystemExit) as ei: + main(["--url", "https://example.com", "--white"]) + assert ei.value.code == 2 + + @patch("rkn_checker.core.check_urls_parallel", return_value=[]) + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + def test_url_in_json_mode_emits_ad_hoc_section( + self, mock_self, mock_parallel, capsys, + ): + main(["--url", "https://example.com", "--json", "--no-self-info"]) + out = capsys.readouterr().out + assert "ad_hoc" in out + # And shouldn't accidentally include the regular sections. + assert "whitelist" not in out + assert "blacklist" not in out + + +class TestIdentifyFlag: + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + @patch("rkn_checker.cli._run_streaming", return_value=([], [])) + def test_identify_propagates_to_streaming(self, mock_stream, mock_self): + main(["--identify"]) + # _run_streaming receives identify as a positional or kwarg — + # check the call args contain True for it. + args, kwargs = mock_stream.call_args + assert True in args or kwargs.get("identify") is True + + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + @patch("rkn_checker.cli._run_streaming", return_value=([], [])) + def test_identify_default_is_off(self, mock_stream, mock_self): + main([]) + args, kwargs = mock_stream.call_args + # The privacy-protective default must be off, not on. + assert kwargs.get("identify", False) is False or False in args + + @patch("rkn_checker.cli.get_self_info", return_value={"ip": "1.2.3.4"}) + @patch("rkn_checker.cli._run_ad_hoc", return_value=[]) + def test_identify_propagates_to_ad_hoc(self, mock_ad_hoc, mock_self): + main(["--url", "https://example.com", "--identify"]) + # _run_ad_hoc(targets, workers, timeout, identify) + args, kwargs = mock_ad_hoc.call_args + # identify is the 4th positional or a kwarg + assert (len(args) >= 4 and args[3] is True) or kwargs.get("identify") is True diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 6ea2b60..1fb0e99 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -10,7 +10,7 @@ from unittest.mock import patch from rkn_checker.core import iter_check_urls -def _slow_check_url(name, url, timeout): +def _slow_check_url(name, url, timeout, identify=False): """Stand-in for check_url that takes a deterministic amount of time based on the target name. 'a' is fastest, 'd' is slowest.""" delays = {"a": 0.05, "b": 0.10, "c": 0.15, "d": 0.20} @@ -78,7 +78,7 @@ class TestExceptionHandling: def test_unexpected_exception_yields_unknown_verdict(self): from rkn_checker.models import CheckResult, Verdict - def _boom(name, url, timeout): + def _boom(name, url, timeout, identify=False): raise RuntimeError("kaboom") urls = {"x": "https://x.example/"}