diffurl.py

#!/usr/bin/env python3
import argparse
import sys
import time
from typing import List, Optional
from pathlib import Path

src_path = Path(__file__).parent / "src"
if src_path.exists():
    sys.path.insert(0, str(src_path))

from diffurl.core import DiffURLOrchestrator, AnalysisRequest, AnalysisMode
from diffurl.norms.unicode_norm import NormalizationForm
from diffurl.variants import VariantType
from diffurl.scoring.tiers import ExploitFlag, SeverityTier
from diffurl.output.reporters import TerminalReporter, JSONReporter, NDJSONReporter, SARIFReporter
from diffurl.output.clipboard import ClipboardManager
from diffurl.psl.updater import PSLUpdater

CLI_VARIANT_MAP = {
    "pct": VariantType.PERCENT_ENCODING,
    "double-pct": VariantType.DOUBLE_PERCENT_ENCODING,
    "hexcase": VariantType.HEX_CASE,
    "backslash": VariantType.BACKSLASH,
    "dotsegs": VariantType.DOT_SEGMENTS,
    "host-forms": VariantType.HOST_FORMS,
    "host_forms": VariantType.HOST_FORMS,
    "dot_segments": VariantType.DOT_SEGMENTS,
}

def _parse_variant_types(spec: str) -> List[VariantType]:
    types: List[VariantType] = []
    for token in (t.strip().lower() for t in spec.split(",") if t.strip()):
        vt = CLI_VARIANT_MAP.get(token)
        if vt:
            types.append(vt)
        else:
            print(f"Warning: Unknown variant type '{token}', skipping", file=sys.stderr)
    if not types:
        types = [
            VariantType.PERCENT_ENCODING,
            VariantType.DOUBLE_PERCENT_ENCODING,
            VariantType.HEX_CASE,
            VariantType.BACKSLASH,
            VariantType.DOT_SEGMENTS,
        ]
    return types

def _parse_norm_forms(spec: str) -> List[NormalizationForm]:
    forms: List[NormalizationForm] = []
    for token in (t.strip() for t in spec.split(",") if t.strip()):
        try:
            forms.append(NormalizationForm(token.upper()))
        except Exception:
            print(f"Warning: Unknown normalization form '{token}', skipping", file=sys.stderr)
    if not forms:
        forms = [NormalizationForm.NFC, NormalizationForm.NFKC]
    return forms

def _normalize_flag_token(token: str) -> Optional[ExploitFlag]:
    t = token.strip().upper()
    if t in ExploitFlag.__members__:
        return ExploitFlag.__members__[t]

    if t == "ETLD1_SHIFT":
        t = "ETLD+1_SHIFT"

    for f in ExploitFlag:
        if f.value.upper() == t:
            return f
    return None


def _split_good_error_results(results):
    good, bad = [], []
    for r in results:
        if getattr(r, "original_analysis", None) is None or ("error" in (r.metadata or {})):
            bad.append(r)
        else:
            good.append(r)
    return good, bad

class DiffURLCli:

    def __init__(self):
        self.parser = self._create_parser()
        self.orchestrator: Optional[DiffURLOrchestrator] = None

    def _create_parser(self) -> argparse.ArgumentParser:
        parser = argparse.ArgumentParser(
            description="diffURL - Terminal URL Canonicalization & Parser-Mismatch Finder",
            epilog="""
Examples:
  # Quick security analysis
  diffurl "https://exɑmple.com/..%2fadmin"

  # Bulk analysis with JSON output
  diffurl urls.txt --json report.json --budget 24

  # CI integration failing on critical issues
  diffurl release_urls.txt --fail-on HOST_SHIFT,ETLD1_SHIFT --json ci_report.json

  # Deep analysis with all techniques
  diffurl target.txt --apex --norms nfc,nfkc,nfd,nfkd --budget 48 --ndjson

  # Update Public Suffix List
  diffurl psl-update
            """,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )

        subparsers = parser.add_subparsers(dest="command")
        sp_psl = subparsers.add_parser("psl-update", help="Update vendored Public Suffix List")
        sp_psl.add_argument("--force", action="store_true", help="Force update even if up-to-date")
        parser.add_argument("urls", nargs="*", help="URLs to analyze (positional)")
        input_group = parser.add_argument_group("Input Options")
        input_group.add_argument("-f", "--file", help="Read URLs from file (one per line)")
        input_group.add_argument("-i", "--stdin", action="store_true", help="Read URLs from standard input")
        analysis_group = parser.add_argument_group("Analysis Modes")
        analysis_group.add_argument("--mode", choices=["standard", "apex", "minimal"], default="standard", help="Analysis mode (default: standard)")
        analysis_group.add_argument("--budget", type=int, default=24, help="Maximum variants per URL (default: 24)")
        analysis_group.add_argument("--apex", action="store_true", help="Use apex set-cover for maximum diversity (alias for --mode apex)")
        norm_group = parser.add_argument_group("Normalization Controls")
        norm_group.add_argument("--norms", default="nfc,nfkc", help="Normalization forms (comma-separated: nfc,nfkc,nfd,nfkd)")
        variant_group = parser.add_argument_group("Variant Generation")
        variant_group.add_argument("--variants", default="pct,double-pct,hexcase,backslash,dotsegs", help="Variant types (comma-separated: pct,double-pct,hexcase,backslash,dotsegs,host-forms)")
        output_group = parser.add_argument_group("Output Formats")
        output_group.add_argument("--json", help="Write JSON report to file")
        output_group.add_argument("--ndjson", action="store_true", help="Stream Newline Delimited JSON to stdout")
        output_group.add_argument("--sarif", help="Write SARIF report for GitHub Code Scanning")
        output_group.add_argument("--minimal", action="store_true", help="Minimal terminal output (summary only)")
        output_group.add_argument("--verbose", action="store_true", help="Detailed terminal output with metadata")
        output_group.add_argument("--copy-poc", action="store_true", help="Copy first minimized witness to clipboard")
        security_group = parser.add_argument_group("Security Policies")
        security_group.add_argument("--fail-on", help="Exit with error if specified flags found (comma-separated)")
        security_group.add_argument("--top-witness-per-severity", type=int, default=1, help="Emit top N minimized witnesses per severity tier")
        data_group = parser.add_argument_group("PSL & Data")
        data_group.add_argument("--psl", choices=["on", "off"], default="on", help="Enable PSL for ETLD+1 analysis (default: on)")
        general_group = parser.add_argument_group("General Options")
        general_group.add_argument("--node-path", default="node", help="Path to Node.js executable (default: node)")
        general_group.add_argument("-v", "--version", action="store_true", help="Show version information")

        return parser

    def parse_args(self, args=None):
        return self.parser.parse_args(args)

    def run(self, args=None) -> int:
        try:
            parsed_args = self.parse_args(args)
            if getattr(parsed_args, "command", None) == "psl-update":
                return self._update_psl(force=getattr(parsed_args, "force", False))

            if parsed_args.version:
                return self._show_version()

            urls = self._get_urls(parsed_args)
            if not urls:
                print("Error: No URLs provided for analysis", file=sys.stderr)
                return 30

            if parsed_args.apex or parsed_args.mode == "apex":
                mode = AnalysisMode.APEX
            elif parsed_args.mode == "minimal":
                mode = AnalysisMode.MINIMAL
            else:
                mode = AnalysisMode.STANDARD

            psl_path = None
            if parsed_args.psl == "off":
                psl_path = str(Path(__file__).parent / "__psl_disabled__")
            self.orchestrator = DiffURLOrchestrator(node_path=parsed_args.node_path, psl_path=psl_path)
            self.orchestrator.initialize()

            if parsed_args.psl == "off":
                try:
                    self.orchestrator._psl_lookup = None  
                    self.orchestrator._initialize_evaluator()
                except Exception:
                    pass

            requests = self._create_analysis_requests(urls, parsed_args, mode)
            results = self._analyze_urls(requests, parsed_args)
            exit_code = self._generate_reports(results, parsed_args)

            if self.orchestrator:
                self.orchestrator.cleanup()

            return exit_code

        except KeyboardInterrupt:
            print("\nAnalysis interrupted by user", file=sys.stderr)
            return 130
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            if self.orchestrator:
                self.orchestrator.cleanup()
            return 1

    def _get_urls(self, args) -> List[str]:
        urls: List[str] = []
        if args.urls:
            urls.extend([u for u in args.urls if u and u != "psl-update"])

        if args.file:
            try:
                with open(args.file, "r", encoding="utf-8") as f:
                    file_urls = [line.strip() for line in f if line.strip()]
                    urls.extend(file_urls)
            except Exception as e:
                raise Exception(f"Failed to read URLs from file: {e}")

        if args.stdin:
            try:
                stdin_urls = [line.strip() for line in sys.stdin if line.strip()]
                urls.extend(stdin_urls)
            except Exception as e:
                raise Exception(f"Failed to read URLs from stdin: {e}")

        seen = set()
        unique_urls = []
        for u in urls:
            if u not in seen:
                seen.add(u)
                unique_urls.append(u)
        return unique_urls

    def _create_analysis_requests(self, urls: List[str], args, mode: AnalysisMode) -> List[AnalysisRequest]:
        norm_forms = _parse_norm_forms(args.norms)
        variant_types = _parse_variant_types(args.variants)

        requests: List[AnalysisRequest] = []
        for url in urls:
            requests.append(
                AnalysisRequest(
                    url=url,
                    mode=mode,
                    normalization_forms=norm_forms,
                    variant_types=variant_types,
                    budget=args.budget,
                    top_witness_per_severity=args.top_witness_per_severity,
                    enable_psl=(args.psl == "on"),
                    enable_confusables=True,
                )
            )
        return requests

    def _analyze_urls(self, requests: List[AnalysisRequest], args) -> List[object]:
        print(f"Analyzing {len(requests)} URL(s) with {args.mode} mode...", file=sys.stderr)

        start_time = time.time()
        health = self.orchestrator.health_check()
        parsers_ok = all(health.get("parsers", {}).values()) if isinstance(health.get("parsers"), dict) else True
        core_ok = health.get("initialized", False) and health.get("normalizer", False) and bool(self.orchestrator._evaluator)

        if not (parsers_ok and core_ok):
            print("Warning: Some components failed health check:", file=sys.stderr)
            for component, status in health.items():
                if component == "parsers" and isinstance(status, dict):
                    for pname, pstatus in status.items():
                        print(f"  parser:{pname}: {'✓' if pstatus else '✗'}", file=sys.stderr)
                else:
                    print(f"  {component}: {'✓' if status else '✗'}", file=sys.stderr)

        if len(requests) == 1:
            result = self.orchestrator.analyze(requests[0])
            results = [result]
        else:
            def progress_callback(completed, total):
                progress = (completed / total) * 100 if total else 100.0
                print(f"Progress: {completed}/{total} ({progress:.1f}%)",
                      end='\r', file=sys.stderr, flush=True)

            results = self.orchestrator.batch_analyze(requests, progress_callback)
            print("", file=sys.stderr)  

        duration = time.time() - start_time
        print(f"Analysis completed in {duration:.2f}s", file=sys.stderr)

        return results

    def _generate_reports(self, results: List[object], args) -> int:
        good_results, error_results = _split_good_error_results(results)

        for er in error_results:
            msg = er.metadata.get("error", "Unknown analysis error")
            target = getattr(er, "request", None)
            which = getattr(target, "url", "unknown")
            print(f"Warning: Failed to analyze '{which}': {msg}", file=sys.stderr)

        fail_on_flags: List[ExploitFlag] = []
        if args.fail_on:
            for token in args.fail_on.split(","):
                flag = _normalize_flag_token(token)
                if flag:
                    fail_on_flags.append(flag)
                else:
                    print(f"Warning: Unknown exploit flag '{token.strip()}', skipping", file=sys.stderr)

        exit_code = 0
        for r in good_results:
            if fail_on_flags and any(f in r.all_flags for f in fail_on_flags):
                exit_code = 10
                break

        produced_output = False

        if args.minimal and not (args.json or args.ndjson or args.sarif):
            self._print_minimal_summary(good_results)
            produced_output = True

        if args.json:
            with open(args.json, "w", encoding="utf-8") as fp:
                jr = JSONReporter(output_stream=fp, pretty=True)
                if len(good_results) == 1:
                    jr.report(good_results[0])
                else:
                    jr.report_batch(good_results)
            produced_output = True

        if args.ndjson:
            ndjr = NDJSONReporter()
            if len(good_results) == 1:
                ndjr.report(good_results[0])
            else:
                ndjr.report_batch(good_results)
            produced_output = True

        if args.sarif:
            with open(args.sarif, "w", encoding="utf-8") as fp:
                sr = SARIFReporter(output_stream=fp)
                if len(good_results) == 1:
                    sr.report(good_results[0])
                else:
                    sr.report_batch(good_results)
            produced_output = True

        if not produced_output and good_results:
            tr = TerminalReporter(verbose=args.verbose)
            if len(good_results) == 1:
                tr.report(good_results[0])
            else:
                tr.report_batch(good_results)

        if args.copy_poc and good_results:
            self._copy_poc_to_clipboard(good_results[0])

        return exit_code

    def _print_minimal_summary(self, results: List[object]) -> None:
        for i, r in enumerate(results, 1):
            sev = r.highest_severity.name if hasattr(r, "highest_severity") else "INFO"
            flags = ",".join(f.value for f in r.all_flags) if hasattr(r, "all_flags") else ""
            print(f"[{i}] {r.request.url}  severity={sev}  findings={len(r.all_flags)}  flags={flags}")

    def _copy_poc_to_clipboard(self, result: object) -> None:
        if not getattr(result, "minimized_witnesses", None):
            print("No minimized witnesses to copy", file=sys.stderr)
            return

        clipboard = ClipboardManager()
        if not clipboard.is_available():
            print("Clipboard not available on this system", file=sys.stderr)
            return

        witness = result.minimized_witnesses[0]
        witness_url = getattr(witness, "minimized_url", str(witness))

        if clipboard.copy_to_clipboard(witness_url):
            print(f"Copied PoC to clipboard: {witness_url}", file=sys.stderr)
        else:
            print("Failed to copy PoC to clipboard", file=sys.stderr)

    def _update_psl(self, force: bool = False) -> int:
        print("Updating Public Suffix List...", file=sys.stderr)
        try:
            updater = PSLUpdater()
            result = updater.update_psl(force=force)
            status = result.get("status")
            if status == "updated":
                print(f"PSL updated successfully: {result.get('new_hash')}", file=sys.stderr)
                return 0
            if status == "skipped":
                print(f"PSL already up to date: {result.get('old_hash')}", file=sys.stderr)
                return 0
            print(f"PSL update failed: {result.get('error', 'Unknown error')}", file=sys.stderr)
            return 40
        except Exception as e:
            print(f"PSL update error: {e}", file=sys.stderr)
            return 40

    def _show_version(self) -> int:
        version_info = {
            "diffURL": "1.0.0",
            "python": sys.version.split()[0],
            "platform": sys.platform,
        }
        print("diffURL - URL Parser Differential Security Tool")
        for key, value in version_info.items():
            print(f"{key}: {value}")
        return 0

def main():
    cli = DiffURLCli()
    exit_code = cli.run()
    sys.exit(exit_code)

if __name__ == "__main__":
    main()