vrr/rkn-block-checker
2
0
Fork 0
mirror of https://github.com/MayersScott/rkn-block-checker.git synced 2026-07-09 17:29:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: support ad-hoc URL probing via --url (can repeat)

Browse source
This commit is contained in:
MayersScott 2026-05-07 14:19:41 +03:00
parent a1f8de5a39
commit 492ce5ea41
3 changed files with 223 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

121
rkn_checker/cli.py
View file

@ -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/<ver>') "
"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())
sys.exit(main())

108
tests/test_cli.py
View file

@ -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

4
tests/test_streaming.py
View file

@ -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/"}