diff --git a/scripts/rugguard.py b/scripts/rugguard.py index dcd51cc..bcf9027 100644 --- a/scripts/rugguard.py +++ b/scripts/rugguard.py @@ -1627,6 +1627,48 @@ def format_markdown(report: RugReport) -> str: def format_json(report: RugReport) -> str: """Format report as pretty JSON.""" return json.dumps(report.to_dict(), indent=2, default=str) + + +def _svg_badge(report: RugReport, style: str = "flat", label: str = "safety") -> str: + """Generate a shields.io-compatible SVG badge.""" + score = report.safety_score + level = report.risk_level + if score >= 70: + bg = "#4c1" + elif score >= 40: + bg = "#e67e22" + elif score >= 20: + bg = "#e74c3c" + else: + bg = "#c0392b" + + label_text = label + value_text = str(score) + "/100 \u2014 " + level + label_w = max(len(label_text) * 7 + 10, 40) + value_w = len(value_text) * 7 + 10 + total_w = label_w + value_w + h = 20 + rx = 3 if style == "flat" else 0 + lx = label_w // 2 + vx = label_w + value_w // 2 + + lines = [] + lines.append('') + lines.append('') + lines.append('') + lines.append('') + lines.append('') + lines.append('') + lines.append('' + _escape_svg(label_text) + '') + lines.append('' + _escape_svg(value_text) + '') + lines.append('') + return "\n".join(lines) + + +def _escape_svg(text: str) -> str: + return (text.replace("&", "&").replace("<", "<").replace(">", ">")) + + # ── CLI Entry Point ──────────────────────────────────────────────────────── def cli_token(args: list[str]) -> None: @@ -1660,12 +1702,39 @@ def cli_wallet(args: list[str]) -> None: if result.get("risky_count", 0) > 0: sys.exit(2) + + +def cli_badge(args: list[str]) -> None: + """Generate an SVG safety score badge for a token.""" + if not args: + print('Usage: python rugguard.py badge [--style flat|flat-square|plastic] [--label TEXT]', file=sys.stderr) + sys.exit(1) + + mint = args[0] + style = "flat" + label = "safety" + + for a in args: + if a.startswith("--style="): + style = a.split("=", 1)[1] + elif a == "--style" and args.index(a) + 1 < len(args): + style = args[args.index(a) + 1] + if a.startswith("--label="): + label = a.split("=", 1)[1] + elif a == "--label" and args.index(a) + 1 < len(args): + label = args[args.index(a) + 1] + + report = rug_check_token(mint.strip()) + print(_svg_badge(report, style=style, label=label)) + + def cli_help() -> None: print("""Solana Rug Guard — On-chain rug-pull detection engine USAGE: python rugguard.py token [--json|--markdown] python rugguard.py wallet + python rugguard.py badge [--style flat|flat-square|plastic] [--label TEXT] python rugguard.py watch [--interval 60] [--iterations 0] [--history PATH] [--webhook URL] [--threshold SCORE] @@ -1685,6 +1754,7 @@ def cli_help() -> None: python rugguard.py token DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263 python rugguard.py token DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263 --markdown python rugguard.py wallet 9WzDXwBbmkg8ZTbNMqUxvQRAyrZzDsGYdLVL9zYtAWWM + python rugguard.py badge DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263 python rugguard.py watch --iterations 1 --threshold 70 ENVIRONMENT: diff --git a/tests/test_checks.py b/tests/test_checks.py index 27ac3ce..4f685fc 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -450,4 +450,73 @@ def test_wallet_scan() -> None: assert "address" in result assert result["address"] == TEST_WALLET assert "total_tokens" in result - assert isinstance(result["total_tokens"], int) + assert isinstance(result['total_tokens'], int) + + +# ── Badge Tests ─────────────────────────────────────────────────────────── + +class TestBadge: + def test_badge_green_low(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=95, risk_level='LOW', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r) + assert '#4c1' in svg # green + assert '95/100' in svg + assert 'LOW' in svg + + def test_badge_yellow_medium(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=55, risk_level='MEDIUM', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r) + assert '#e67e22' in svg # yellow + assert '55/100' in svg + + def test_badge_red_high(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=25, risk_level='HIGH', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r) + assert '#e74c3c' in svg # red + + def test_badge_darkred_critical(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=10, risk_level='CRITICAL', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r) + assert '#c0392b' in svg # dark red + + def test_badge_custom_label(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=80, risk_level='LOW', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r, label='rugcheck') + assert 'rugcheck' in svg + + def test_badge_is_valid_svg(self): + from rugguard import _svg_badge, RugScore, RugFlags, TokenMeta + flags = RugFlags() + r = RugReport( + token=TokenMeta(address='A'), safety_score=80, risk_level='LOW', + score=RugScore(), flags=flags, warnings=[], recommendation='', + ) + svg = _svg_badge(r) + assert svg.startswith('') + assert 'xmlns=' in svg