Add batch processing, html reports, and config

Author: Black-Coffee-Ramen

herald/__init__.py

@@ -1,4 +1,4 @@
-"""
+"""
 HERALD — API-free phishing domain detection for critical infrastructure.
 
 Usage:
@@ -7,5 +7,5 @@
     result = predictor.predict("sbi-login-secure.xyz")
 """
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 __author__ = "Athiyo Chakma"

herald/cli.py

@@ -12,6 +12,7 @@
 from rich.panel import Panel
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.table import Table
+from rich.prompt import Confirm
 
 from herald.core.security import SSRFProtectionError
 from herald.investigation.persistence import find_report
@@ -23,7 +24,16 @@
 
 def main(argv: list[str] | None = None) -> int:
     parser = build_parser()
-    args = parser.parse_args(argv)
+    args, _ = parser.parse_known_args(argv)
+
+    if args.version:
+        console.print(Panel.fit(
+            f"[bold cyan]HERALD[/bold cyan] v{__version__}\n"
+            "[dim]Heuristic & Ensemble Risk Assessment for Lookalike Domains[/dim]",
+            border_style="cyan",
+            box=box.ROUNDED
+        ))
+        return 0
 
     if args.command == "investigate":
         return run_investigate(args)
@@ -33,14 +43,20 @@ def main(argv: list[str] | None = None) -> int:
         return run_screenshot(args)
     if args.command == "report":
         return run_report(args)
+    if args.command == "update":
+        return run_update(args)
+    if args.command == "config":
+        return run_config(args)
+    if args.command == "cleanup":
+        return run_cleanup(args)
 
     parser.print_help()
     return 1
 
 
 def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(prog="herald", description="HERALD phishing investigation CLI")
-    parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")
+    parser.add_argument("-v", "--version", action="store_true", help="show program's version number and exit")
     subparsers = parser.add_subparsers(dest="command")
 
     investigate = subparsers.add_parser("investigate", help="Run a full URL investigation")
@@ -64,18 +80,80 @@ def build_parser() -> argparse.ArgumentParser:
     report = subparsers.add_parser("report", help="Show a persisted investigation by trace ID")
     report.add_argument("trace_id")
     report.add_argument("--json", action="store_true", help="Print machine-readable JSON")
+    report.add_argument("--open", action="store_true", help="Open the evidence directory in the file explorer")
+
+    update = subparsers.add_parser("update", help="Update HERALD Investigator to the latest version")
+
+    config = subparsers.add_parser("config", help="Manage HERALD configuration")
+    config_sub = config.add_subparsers(dest="config_command", required=True)
+    config_sub.add_parser("show", help="Show current configuration")
+    config_set = config_sub.add_parser("set", help="Set a configuration value")
+    config_set.add_argument("key")
+    config_set.add_argument("value")
+
+    cleanup = subparsers.add_parser("cleanup", help="Clean up old evidence directories")
+    cleanup.add_argument("--older-than", type=int, default=30, help="Delete investigations older than N days (default: 30)")
 
     return parser
 
 
 def run_investigate(args: argparse.Namespace) -> int:
+    import os
+    from herald import config
+    
+    if os.path.isfile(args.target):
+        return run_batch_investigate(args)
+        
     pipeline = InvestigationPipeline(scorer_type=args.scorer)
+    
+    do_visual = not args.no_visual
+    if "screenshot" in config.get_config() and not args.no_visual:
+        do_visual = config.get("screenshot")
+        
     try:
-        result = execute_pipeline(pipeline, args.target, include_visual=not args.no_visual, allow_private=args.allow_private, quiet=args.json)
+        result = execute_pipeline(pipeline, args.target, include_visual=do_visual, allow_private=args.allow_private, quiet=args.json)
     except SSRFProtectionError as exc:
         return emit_ssrf_error(exc, args.target, "investigate", as_json=args.json)
     return emit_result(result.to_dict(), as_json=args.json, explain=args.explain)
 
+def run_batch_investigate(args: argparse.Namespace) -> int:
+    from collections import Counter
+    pipeline = InvestigationPipeline(scorer_type=args.scorer)
+    pipeline.engine.preload()
+    
+    with open(args.target, "r", encoding="utf-8") as f:
+        targets = [line.strip() for line in f if line.strip() and not line.startswith("#")]
+        
+    if not targets:
+        console.print(f"[red]No targets found in {args.target}[/red]")
+        return 1
+
+    console.print(f"[bold cyan]Starting batch investigation of {len(targets)} targets...[/bold cyan]")
+    
+    stats = Counter()
+    
+    with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+        task = progress.add_task(f"Processing 0/{len(targets)}...", total=len(targets))
+        
+        for i, target in enumerate(targets, 1):
+            progress.update(task, description=f"Processing {i}/{len(targets)}: {target}")
+            try:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+                    res = pipeline.investigate(target, include_visual=not args.no_visual, allow_private=args.allow_private)
+                stats[res.verdict] += 1
+            except Exception:
+                stats["Errors"] += 1
+
+    console.print("\n[bold]Batch Investigation Summary[/bold]")
+    console.print(f"Total processed: {len(targets)}")
+    console.print(f"Phishing (High Risk): [red]{stats.get('Phishing', 0)}[/red]")
+    console.print(f"Suspected: [yellow]{stats.get('Suspected', 0)}[/yellow]")
+    console.print(f"Clean: [green]{stats.get('Clean', 0)}[/green]")
+    if stats.get("Errors"):
+        console.print(f"Errors: [bright_black]{stats['Errors']}[/bright_black]")
+        
+    return 0
+
 
 def run_analyze(args: argparse.Namespace) -> int:
     pipeline = InvestigationPipeline()
@@ -111,9 +189,97 @@ def run_report(args: argparse.Namespace) -> int:
     if not report:
         console.print(f"[red]No persisted report found for trace ID[/red] {args.trace_id}")
         return 1
+        
+    if args.open:
+        import subprocess
+        import os
+        evidence_dir = report.get("evidence_dir")
+        if evidence_dir and os.path.exists(evidence_dir):
+            if sys.platform == "win32":
+                os.startfile(evidence_dir)
+            elif sys.platform == "darwin":
+                subprocess.Popen(["open", evidence_dir])
+            else:
+                subprocess.Popen(["xdg-open", evidence_dir])
+            console.print(f"[green]Opened evidence directory:[/green] {evidence_dir}")
+            return 0
+            
     return emit_result(report, as_json=args.json)
 
 
+def run_update(args: argparse.Namespace) -> int:
+    import urllib.request
+    import subprocess
+    
+    console.print("[cyan]Checking for HERALD Investigator updates...[/cyan]")
+    try:
+        req = urllib.request.Request("https://pypi.org/pypi/herald-investigator/json")
+        with urllib.request.urlopen(req, timeout=5) as response:
+            data = json.loads(response.read().decode())
+            latest_version = data["info"]["version"]
+    except Exception as e:
+        console.print(f"[red]Failed to check PyPI for updates:[/red] {e}")
+        return 1
+
+    console.print(f"Current Version: [bold]{__version__}[/bold]")
+    console.print(f"Latest Version:  [bold green]{latest_version}[/bold green]\n")
+
+    # Basic string comparison (adequate for simple semver)
+    if __version__ == latest_version:
+        console.print("You are already on the latest version!")
+        return 0
+
+    if Confirm.ask("Would you like to install the update now?"):
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "herald-investigator"])
+            console.print("\n[bold green]✅ Update successful![/bold green]")
+            return 0
+        except subprocess.CalledProcessError as e:
+            console.print(f"\n[bold red]❌ Update failed:[/bold red] {e}")
+            return 1
+    return 0
+
+
+def run_config(args: argparse.Namespace) -> int:
+    from herald import config
+    if args.config_command == "show":
+        c = config.get_config()
+        if not c:
+            console.print("Configuration is empty.")
+        else:
+            for k, v in c.items():
+                console.print(f"[cyan]{k}[/cyan]: {v}")
+    elif args.config_command == "set":
+        config.set_val(args.key, args.value)
+        console.print(f"[green]Set[/green] {args.key} = {args.value}")
+    return 0
+
+
+def run_cleanup(args: argparse.Namespace) -> int:
+    import time
+    import shutil
+    from pathlib import Path
+    
+    evidence_dir = Path("evidence")
+    if not evidence_dir.exists():
+        console.print("No evidence directory found.")
+        return 0
+        
+    cutoff = time.time() - (args.older_than * 86400)
+    count = 0
+    
+    for path in evidence_dir.iterdir():
+        if path.is_dir():
+            json_file = path / "investigation.json"
+            check_path = json_file if json_file.exists() else path
+            if check_path.stat().st_mtime < cutoff:
+                shutil.rmtree(path)
+                count += 1
+                
+    console.print(f"[green]Cleanup complete.[/green] Deleted {count} old investigations.")
+    return 0
+
+
 def emit_result(result: dict, *, as_json: bool, explain: bool = False) -> int:
     if as_json:
         console.print_json(json.dumps(result))

herald/investigation/persistence.py

@@ -19,10 +19,12 @@ def save_investigation(result: dict[str, Any], evidence_dir: str) -> dict[str, s
     base.mkdir(parents=True, exist_ok=True)
     json_path = base / "investigation.json"
     md_path = base / "report.md"
-    result.setdefault("visual", {}).setdefault("report_paths", {"json": os.fspath(json_path), "markdown": os.fspath(md_path)})
+    html_path = base / "report.html"
+    result.setdefault("visual", {}).setdefault("report_paths", {"json": os.fspath(json_path), "markdown": os.fspath(md_path), "html": os.fspath(html_path)})
 
     json_path.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
     md_path.write_text(render_markdown(result), encoding="utf-8")
+    html_path.write_text(render_html(result), encoding="utf-8")
 
     latest_dir = Path("evidence")
     latest_dir.mkdir(exist_ok=True)
@@ -37,7 +39,7 @@ def save_investigation(result: dict[str, Any], evidence_dir: str) -> dict[str, s
             "completed_at": result["completed_at"],
         }, sort_keys=True) + "\n")
 
-    return {"json": os.fspath(json_path), "markdown": os.fspath(md_path)}
+    return {"json": os.fspath(json_path), "markdown": os.fspath(md_path), "html": os.fspath(html_path)}
 
 
 def find_report(trace_id: str, root: str = "evidence") -> dict[str, Any] | None:
@@ -137,3 +139,109 @@ def render_markdown(result: dict[str, Any]) -> str:
             lines.append(f"- {error}")
 
     return "\n".join(lines) + "\n"
+
+def render_html(result: dict[str, Any]) -> str:
+    """Generate a clean, professional HTML report from the investigation result."""
+    dns = result.get("dns", {})
+    tls = result.get("tls", {})
+    visual = result.get("visual", {})
+    summary = result.get("summary", {})
+    risk_factors = result.get("risk_factors") or []
+    
+    # Base styling
+    html = [
+        "<!DOCTYPE html>",
+        "<html lang='en'>",
+        "<head>",
+        "    <meta charset='UTF-8'>",
+        f"    <title>HERALD Report: {result.get('domain', 'Unknown')}</title>",
+        "    <style>",
+        "        body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif; line-height: 1.6; color: #333; max-width: 1000px; margin: 0 auto; padding: 2rem; background-color: #f9f9fb; }",
+        "        h1, h2, h3 { color: #1a1a2e; }",
+        "        .header { background-color: #ffffff; padding: 2rem; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); margin-bottom: 2rem; border-top: 5px solid #2ecc71; }",
+        "        .header.Phishing { border-top-color: #e74c3c; }",
+        "        .header.Suspected { border-top-color: #f1c40f; }",
+        "        .section { background-color: #ffffff; padding: 1.5rem; border-radius: 8px; box-shadow: 0 1px 3px rgba(0,0,0,0.05); margin-bottom: 1.5rem; }",
+        "        table { width: 100%; border-collapse: collapse; margin-top: 1rem; }",
+        "        th, td { padding: 12px; text-align: left; border-bottom: 1px solid #eaeaea; }",
+        "        th { background-color: #f8f9fa; font-weight: 600; color: #555; }",
+        "        .badge { display: inline-block; padding: 0.25em 0.6em; font-size: 0.85em; font-weight: 700; line-height: 1; text-align: center; white-space: nowrap; vertical-align: baseline; border-radius: 0.25rem; }",
+        "        .badge.high { background-color: #fee2e2; color: #991b1b; }",
+        "        .badge.medium { background-color: #fef3c7; color: #92400e; }",
+        "        .badge.low { background-color: #f3f4f6; color: #374151; }",
+        "        .badge.info { background-color: #e0f2fe; color: #075985; }",
+        "        .screenshot-container { max-width: 100%; margin-top: 1rem; border: 1px solid #eee; border-radius: 4px; overflow: hidden; }",
+        "        .screenshot-container img { max-width: 100%; height: auto; display: block; }",
+        "    </style>",
+        "</head>",
+        "<body>",
+    ]
+    
+    headline = summary.get('headline') or f"{result.get('domain')} is {result.get('verdict')} with score {result.get('phishing_score')}"
+    
+    html.extend([
+        f"    <div class='header {result.get('verdict', 'Clean')}'>",
+        "        <h1>HERALD Investigation Report</h1>",
+        f"        <p style='font-size: 1.2rem;'><strong>{headline}</strong></p>",
+        "        <table style='width: 100%; background: #fdfdfd;'>",
+        f"            <tr><td><strong>Trace ID:</strong></td><td><code>{result.get('trace_id')}</code></td><td><strong>Verdict:</strong></td><td><strong>{result.get('verdict')}</strong></td></tr>",
+        f"            <tr><td><strong>Domain:</strong></td><td><code>{result.get('domain')}</code></td><td><strong>Score:</strong></td><td><code>{result.get('phishing_score')}</code></td></tr>",
+        f"            <tr><td><strong>Completed:</strong></td><td>{result.get('completed_at')}</td><td><strong>Target:</strong></td><td><code>{result.get('url')}</code></td></tr>",
+        "        </table>",
+        "    </div>",
+    ])
+
+    # Why Flagged
+    html.append("<div class='section'><h2>Why Flagged</h2><ul>")
+    for reason in summary.get("why_flagged", ["No strong phishing indicators were observed."]):
+        html.append(f"<li>{reason}</li>")
+    html.append("</ul></div>")
+
+    # Risk Factors
+    html.append("<div class='section'><h2>Risk Factors</h2><table>")
+    html.append("<tr><th>Factor</th><th>Severity</th><th>Impact</th><th>Explanation</th></tr>")
+    if risk_factors:
+        for factor in risk_factors:
+            sev = factor.get('severity', 'info')
+            html.append(f"<tr><td>{factor.get('name', 'Unknown')}</td><td><span class='badge {sev}'>{sev}</span></td><td>{factor.get('score_impact', 0)}</td><td>{factor.get('detail', '')}</td></tr>")
+    else:
+        html.append("<tr><td colspan='4'>No notable risk factors recorded.</td></tr>")
+    html.append("</table></div>")
+
+    # DNS and TLS
+    html.append("<div style='display: flex; gap: 1.5rem;'>")
+    
+    # DNS
+    html.append("<div class='section' style='flex: 1;'><h2>DNS & WHOIS</h2><table>")
+    html.append(f"<tr><th>Registrar</th><td>{dns.get('registrar') or 'unknown'}</td></tr>")
+    html.append(f"<tr><th>Domain Age</th><td>{dns.get('domain_age_days') if dns.get('domain_age_days') is not None else 'unknown'} days</td></tr>")
+    html.append(f"<tr><th>A Records</th><td>{', '.join(dns.get('a_records') or []) or 'none'}</td></tr>")
+    html.append(f"<tr><th>MX Records</th><td>{', '.join(dns.get('mx_records') or []) or 'none'}</td></tr>")
+    html.append("</table></div>")
+
+    # TLS
+    html.append("<div class='section' style='flex: 1;'><h2>TLS Certificate</h2><table>")
+    html.append(f"<tr><th>Has TLS</th><td>{tls.get('has_tls', False)}</td></tr>")
+    html.append(f"<tr><th>Issuer</th><td>{tls.get('issuer') or 'unknown'}</td></tr>")
+    html.append(f"<tr><th>SAN Match</th><td>{tls.get('domain_in_san', False)}</td></tr>")
+    html.append(f"<tr><th>Days Remaining</th><td>{tls.get('days_remaining') if tls.get('days_remaining') is not None else 'unknown'}</td></tr>")
+    html.append("</table></div>")
+    
+    html.append("</div>") # End flex
+
+    # Evidence / Screenshot
+    html.append("<div class='section'><h2>Visual Evidence</h2>")
+    ocr = ", ".join((visual.get('ocr_findings') or {}).get('phrases_found') or []) or "none"
+    html.append(f"<p><strong>OCR Suspicious Phrases:</strong> {ocr}</p>")
+    
+    screenshot_path = visual.get('screenshot_path')
+    if screenshot_path:
+        # Resolve path relative to the HTML report (which is in the same dir as the screenshots folder)
+        rel_path = "screenshots/homepage.png" if "homepage.png" in screenshot_path else screenshot_path.split("/")[-1]
+        html.append(f"<div class='screenshot-container'><img src='{rel_path}' alt='Screenshot of {result.get('domain')}'></div>")
+    else:
+        html.append("<p><em>Screenshot not captured.</em></p>")
+    html.append("</div>")
+
+    html.append("</body></html>")
+    return "\n".join(html)

herald/predict_with_fallback.py

@@ -64,13 +64,15 @@ class PhishingPredictorV3:
     Uses Lexical + Network (SSL/DNS) + WHOIS features with a two-stage pipeline.
     Maintains compatibility with v3 class name for system-wide integration.
     """
-    def __init__(self, model_path="models/ensemble_v7.joblib", config_path="config.yaml"):
+    def __init__(self, model_path=None, config_path=None):
+        base_dir = os.path.dirname(os.path.abspath(__file__))
+        if model_path is None:
+            model_path = os.path.join(base_dir, "models", "ensemble_v7.joblib")
+        if config_path is None:
+            config_path = os.path.join(base_dir, "config.yaml")
+
         print(f"Loading HERALD v7 Phishing Predictor from {model_path}...")
-        if not os.path.exists(model_path):
-            # Fallback to absolute path or project root if needed
-            base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-            model_path = os.path.join(base_dir, model_path)
-            
+        
         if not os.path.exists(model_path):
             raise FileNotFoundError(f"Model file {model_path} not found.")