httpprobe.py

#!/usr/bin/env python3
"""
httpprobe -- Check HTTP endpoints fast. Batch URL health checks, zero deps.

Verify endpoints are up, check response times, validate JSON, inspect SSL.
Perfect for deployment smoke tests, monitoring, and CI/CD pipelines.

Usage:
    py httpprobe.py https://api.example.com/health                    # Single URL
    py httpprobe.py url1 url2 url3                                    # Multiple URLs
    py httpprobe.py -f urls.txt                                       # URLs from file
    py httpprobe.py https://api.example.com -H "Auth: Bearer tok"     # Custom header
    py httpprobe.py https://api.example.com --json-field status ok     # Validate JSON
    py httpprobe.py https://example.com --ssl                         # Show SSL info
    py httpprobe.py url1 url2 --parallel 5 --timeout 10               # Parallel + timeout
    py httpprobe.py url1 url2 --output json                           # JSON output
    py httpprobe.py url1 --expect 200                                 # Assert status code
"""

import argparse
import json
import os
import ssl
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

# Colors
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
RED = "\033[31m"
MAGENTA = "\033[35m"


def color_supported() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    if sys.platform == "win32":
        return bool(os.environ.get("TERM") or os.environ.get("WT_SESSION"))
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


USE_COLOR = color_supported()


def c(code: str, text: str) -> str:
    return f"{code}{text}{RESET}" if USE_COLOR else text


def status_color(code: int) -> str:
    if 200 <= code < 300:
        return GREEN
    elif 300 <= code < 400:
        return YELLOW
    else:
        return RED


def format_ms(seconds: float) -> str:
    if seconds < 1.0:
        return f"{seconds * 1000:.0f}ms"
    return f"{seconds:.2f}s"


def get_ssl_info(hostname: str, port: int = 443) -> dict:
    """Get SSL certificate information."""
    try:
        ctx = ssl.create_default_context()
        with ctx.wrap_socket(
            __import__("socket").create_connection((hostname, port), timeout=5),
            server_hostname=hostname
        ) as sock:
            cert = sock.getpeercert()
            if not cert:
                return {"error": "No certificate"}

            subject = dict(x[0] for x in cert.get("subject", ()))
            issuer = dict(x[0] for x in cert.get("issuer", ()))

            not_after = cert.get("notAfter", "")
            try:
                expiry = datetime.strptime(not_after, "%b %d %H:%M:%S %Y %Z")
                days_left = (expiry - datetime.utcnow()).days
            except (ValueError, TypeError):
                expiry = None
                days_left = None

            return {
                "common_name": subject.get("commonName", ""),
                "issuer": issuer.get("organizationName", issuer.get("commonName", "")),
                "expires": not_after,
                "days_left": days_left,
                "san": [e[1] for e in cert.get("subjectAltName", ())],
                "serial": cert.get("serialNumber", ""),
            }
    except Exception as e:
        return {"error": str(e)}


def probe_url(url: str, method: str = "GET", headers: dict = None,
              body: str = None, timeout: float = 10,
              follow_redirects: bool = True, show_ssl: bool = False,
              json_field: tuple = None, expect_status: int = None) -> dict:
    """Probe a single URL and return results."""
    result = {
        "url": url,
        "method": method,
        "status": None,
        "time_s": None,
        "size_bytes": None,
        "content_type": None,
        "error": None,
        "ok": False,
        "headers": {},
        "ssl": None,
        "json_check": None,
    }

    # Ensure URL has scheme
    if not url.startswith(("http://", "https://")):
        url = "https://" + url
        result["url"] = url

    req = Request(url, method=method)
    if headers:
        for k, v in headers.items():
            req.add_header(k, v)
    if body:
        req.data = body.encode("utf-8")
        if "Content-Type" not in (headers or {}):
            req.add_header("Content-Type", "application/json")

    req.add_header("User-Agent", "httpprobe/1.0")

    # SSL context
    ctx = ssl.create_default_context()

    start = time.perf_counter()
    try:
        resp = urlopen(req, timeout=timeout, context=ctx)
        elapsed = time.perf_counter() - start

        result["status"] = resp.status
        result["time_s"] = elapsed
        result["content_type"] = resp.headers.get("Content-Type", "")
        result["headers"] = dict(resp.headers)

        body_bytes = resp.read()
        result["size_bytes"] = len(body_bytes)
        result["ok"] = True

        # JSON field validation
        if json_field:
            try:
                data = json.loads(body_bytes.decode("utf-8"))
                field_name, expected_value = json_field
                actual = _extract_field(data, field_name)
                matched = str(actual) == expected_value
                result["json_check"] = {
                    "field": field_name,
                    "expected": expected_value,
                    "actual": actual,
                    "match": matched,
                }
                if not matched:
                    result["ok"] = False
            except (json.JSONDecodeError, Exception) as e:
                result["json_check"] = {"error": str(e)}
                result["ok"] = False

    except HTTPError as e:
        elapsed = time.perf_counter() - start
        result["status"] = e.code
        result["time_s"] = elapsed
        result["error"] = str(e.reason)
        result["ok"] = 200 <= e.code < 400

    except URLError as e:
        elapsed = time.perf_counter() - start
        result["time_s"] = elapsed
        result["error"] = str(e.reason)

    except Exception as e:
        elapsed = time.perf_counter() - start
        result["time_s"] = elapsed
        result["error"] = str(e)

    # Expected status check
    if expect_status and result["status"] != expect_status:
        result["ok"] = False

    # SSL info
    if show_ssl and url.startswith("https://"):
        try:
            from urllib.parse import urlparse
            parsed = urlparse(url)
            hostname = parsed.hostname
            port = parsed.port or 443
            result["ssl"] = get_ssl_info(hostname, port)
        except Exception as e:
            result["ssl"] = {"error": str(e)}

    return result


def _extract_field(data, field_path: str):
    """Extract a field from nested data using dot notation."""
    parts = field_path.split(".")
    current = data
    for part in parts:
        if isinstance(current, dict):
            current = current.get(part)
        elif isinstance(current, list):
            try:
                current = current[int(part)]
            except (ValueError, IndexError):
                return None
        else:
            return None
    return current


def print_result(r: dict, verbose: bool = False):
    """Print a single probe result."""
    status = r["status"]
    url = r["url"]
    elapsed = r["time_s"] or 0
    size = r["size_bytes"]

    # Status indicator
    if r["ok"]:
        icon = c(GREEN, "OK")
    elif r["error"]:
        icon = c(RED, "FAIL")
    else:
        icon = c(YELLOW, "WARN")

    # Status code
    if status:
        sc = c(status_color(status), str(status))
    else:
        sc = c(RED, "---")

    # Size
    if size is not None:
        if size < 1024:
            size_str = f"{size}B"
        elif size < 1024 * 1024:
            size_str = f"{size / 1024:.1f}KB"
        else:
            size_str = f"{size / (1024 * 1024):.1f}MB"
    else:
        size_str = "---"

    time_str = format_ms(elapsed)

    print(f"  {icon}  {sc}  {time_str:>7}  {size_str:>8}  {url}")

    if r["error"]:
        print(c(RED, f"         Error: {r['error']}"))

    if r.get("json_check"):
        jc = r["json_check"]
        if "error" in jc:
            print(c(RED, f"         JSON check error: {jc['error']}"))
        elif jc.get("match"):
            print(c(GREEN, f"         JSON: {jc['field']} = {jc['actual']} (matched)"))
        else:
            print(c(RED, f"         JSON: {jc['field']} = {jc['actual']} (expected: {jc['expected']})"))

    if r.get("ssl"):
        si = r["ssl"]
        if "error" in si:
            print(c(RED, f"         SSL: {si['error']}"))
        else:
            days = si.get("days_left")
            if days is not None:
                if days < 30:
                    day_color = RED
                elif days < 90:
                    day_color = YELLOW
                else:
                    day_color = GREEN
                print(f"         SSL: {si['issuer']} | expires {si['expires']} ({c(day_color, f'{days}d left')})")
            else:
                print(f"         SSL: {si['issuer']} | expires {si['expires']}")

    if verbose and r.get("content_type"):
        print(c(DIM, f"         Type: {r['content_type']}"))


def print_summary(results: list[dict]):
    """Print a summary of all results."""
    total = len(results)
    ok_count = sum(1 for r in results if r["ok"])
    fail_count = total - ok_count
    times = [r["time_s"] for r in results if r["time_s"]]
    avg_time = sum(times) / len(times) if times else 0

    print()
    if fail_count == 0:
        print(f"  {c(GREEN, 'All passed')} ({total}/{total}) | avg {format_ms(avg_time)}")
    else:
        print(f"  {c(GREEN, f'{ok_count} passed')}, {c(RED, f'{fail_count} failed')} ({total} total) | avg {format_ms(avg_time)}")
    print()


def main():
    parser = argparse.ArgumentParser(
        description="httpprobe -- check HTTP endpoints fast",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("urls", nargs="*", help="URL(s) to probe")
    parser.add_argument("-f", "--file", help="Read URLs from file (one per line)")
    parser.add_argument("-m", "--method", default="GET", help="HTTP method (default: GET)")
    parser.add_argument("-H", "--header", action="append", default=[], help="Header (Key: Value). Repeatable.")
    parser.add_argument("-d", "--data", help="Request body (for POST/PUT)")
    parser.add_argument("-t", "--timeout", type=float, default=10, help="Timeout in seconds (default: 10)")
    parser.add_argument("-p", "--parallel", type=int, default=4, help="Max parallel requests (default: 4)")
    parser.add_argument("--ssl", action="store_true", help="Show SSL certificate info")
    parser.add_argument("--json-field", nargs=2, metavar=("FIELD", "VALUE"),
                        help="Validate JSON response field (dot notation)")
    parser.add_argument("--expect", type=int, metavar="CODE", help="Expected HTTP status code")
    parser.add_argument("--output", choices=["text", "json"], default="text", help="Output format")
    parser.add_argument("-v", "--verbose", action="store_true", help="Show extra details")
    parser.add_argument("--no-follow", action="store_true", help="Don't follow redirects")

    args = parser.parse_args()

    # Collect URLs
    urls = list(args.urls or [])
    if args.file:
        try:
            with open(args.file, "r", encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if line and not line.startswith("#"):
                        urls.append(line)
        except FileNotFoundError:
            print(f"Error: file not found: {args.file}", file=sys.stderr)
            sys.exit(1)

    if not urls:
        parser.print_help()
        sys.exit(1)

    # Parse headers
    headers = {}
    for h in args.header:
        if ":" in h:
            k, v = h.split(":", 1)
            headers[k.strip()] = v.strip()

    json_field = tuple(args.json_field) if args.json_field else None

    if args.output == "text":
        print()
        print(c(BOLD, "  httpprobe"))
        print(c(DIM, f"  Checking {len(urls)} endpoint(s)..."))
        print()

    # Probe URLs in parallel
    results = []
    with ThreadPoolExecutor(max_workers=args.parallel) as pool:
        futures = {}
        for url in urls:
            future = pool.submit(
                probe_url, url,
                method=args.method,
                headers=headers if headers else None,
                body=args.data,
                timeout=args.timeout,
                follow_redirects=not args.no_follow,
                show_ssl=args.ssl,
                json_field=json_field,
                expect_status=args.expect,
            )
            futures[future] = url

        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            if args.output == "text":
                print_result(result, verbose=args.verbose)

    # Sort by original URL order
    url_order = {url: i for i, url in enumerate(urls)}
    results.sort(key=lambda r: url_order.get(r["url"], 999))

    if args.output == "json":
        clean = []
        for r in results:
            entry = {
                "url": r["url"],
                "method": r["method"],
                "status": r["status"],
                "time_ms": round(r["time_s"] * 1000, 1) if r["time_s"] else None,
                "size_bytes": r["size_bytes"],
                "ok": r["ok"],
                "error": r["error"],
            }
            if r.get("ssl"):
                entry["ssl"] = r["ssl"]
            if r.get("json_check"):
                entry["json_check"] = r["json_check"]
            clean.append(entry)
        print(json.dumps(clean, indent=2))
    else:
        print_summary(results)

    # Exit code: 0 if all ok, 1 if any failed
    if not all(r["ok"] for r in results):
        sys.exit(1)


if __name__ == "__main__":
    main()