path: root/src/cvedb/cli.py

import json
import sys
from enum import Enum

import typer
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text

from cvedb.client import CVEDBClient
from cvedb.models import CVE, CVEWithCPEs

CONTEXT = {"help_option_names": ["-h", "--help"]}
app = typer.Typer(help="CLI client for Shodan's CVEDB API", no_args_is_help=True, context_settings=CONTEXT)
console = Console()
err_console = Console(stderr=True)


class Format(str, Enum):
    table = "table"
    json = "json"
    tsv = "tsv"


def get_format(fmt: Format | None) -> Format:
    if fmt:
        return fmt
    return Format.table if sys.stdout.isatty() else Format.tsv


def severity_color(cvss: float | None) -> str:
    if cvss is None:
        return "dim"
    if cvss >= 9.0:
        return "red bold"
    if cvss >= 7.0:
        return "orange1"
    if cvss >= 4.0:
        return "yellow"
    return "green"


def format_cvss(cvss: float | None) -> str:
    return f"{cvss:.1f}" if cvss is not None else "-"


def cve_to_dict(cve: CVE | CVEWithCPEs) -> dict:
    data = cve.model_dump(mode="json")
    return data


def render_cve_detail(cve: CVEWithCPEs, fmt: Format):
    if fmt == Format.json:
        print(json.dumps(cve_to_dict(cve), indent=2))
        return

    if fmt == Format.tsv:
        score = cve.cvss_v3 or cve.cvss_v2 or cve.cvss
        fields = [
            cve.cve_id,
            format_cvss(score),
            f"{cve.epss:.4f}" if cve.epss else "",
            "YES" if cve.kev else "no",
            cve.published_time.strftime("%Y-%m-%d") if cve.published_time else "",
            (cve.summary or "").replace("\t", " ").replace("\n", " "),
        ]
        print("\t".join(fields))
        return

    severity = severity_color(cve.cvss_v3 or cve.cvss_v2 or cve.cvss)
    score = cve.cvss_v3 or cve.cvss_v2 or cve.cvss

    title = Text()
    title.append(cve.cve_id, style="bold cyan")
    if cve.kev:
        title.append(" [KEV]", style="red bold")
    if cve.ransomware_campaign:
        title.append(" [RANSOMWARE]", style="magenta bold")

    content = Text()
    content.append("CVSS: ", style="bold")
    content.append(format_cvss(score), style=severity)
    if cve.cvss_v3:
        content.append(f" (v3: {cve.cvss_v3})")
    if cve.cvss_v2:
        content.append(f" (v2: {cve.cvss_v2})")
    content.append("\n")

    if cve.epss is not None:
        content.append("EPSS: ", style="bold")
        content.append(f"{cve.epss:.4f}")
        if cve.ranking_epss is not None:
            content.append(f" (top {cve.ranking_epss*100:.1f}%)")
        content.append("\n")

    if cve.published_time:
        content.append("Published: ", style="bold")
        content.append(cve.published_time.strftime("%Y-%m-%d"))
        content.append("\n")

    if cve.propose_action:
        content.append("Action: ", style="bold")
        content.append(cve.propose_action, style="yellow")
        content.append("\n")

    content.append("\n")
    content.append(cve.summary or "No description available.", style="italic")

    console.print(Panel(content, title=title, border_style=severity))

    if cve.cpes:
        console.print("\n[bold]Affected Products:[/]")
        for cpe in cve.cpes[:20]:
            console.print(f"  [dim]{cpe}[/]")
        if len(cve.cpes) > 20:
            console.print(f"  [dim]... and {len(cve.cpes) - 20} more[/]")

    if cve.references:
        console.print("\n[bold]References:[/]")
        for ref in cve.references[:10]:
            console.print(f"  [link={ref}]{ref}[/link]")
        if len(cve.references) > 10:
            console.print(f"  [dim]... and {len(cve.references) - 10} more[/]")


def render_cve_table(cves: list[CVE], fmt: Format):
    if fmt == Format.json:
        print(json.dumps([cve_to_dict(c) for c in cves], indent=2))
        return

    if fmt == Format.tsv:
        print("cve_id\tcvss\tepss\tkev\tpublished\tsummary")
        for cve in cves:
            score = cve.cvss_v3 or cve.cvss_v2 or cve.cvss
            fields = [
                cve.cve_id,
                format_cvss(score),
                f"{cve.epss:.4f}" if cve.epss else "",
                "YES" if cve.kev else "no",
                cve.published_time.strftime("%Y-%m-%d") if cve.published_time else "",
                (cve.summary or "").replace("\t", " ").replace("\n", " ")[:100],
            ]
            print("\t".join(fields))
        return

    table = Table(show_header=True, header_style="bold")
    table.add_column("CVE ID", style="cyan")
    table.add_column("CVSS", justify="right")
    table.add_column("EPSS", justify="right")
    table.add_column("KEV", justify="center")
    table.add_column("Published")
    table.add_column("Summary", max_width=50, overflow="ellipsis")

    for cve in cves:
        score = cve.cvss_v3 or cve.cvss_v2 or cve.cvss
        table.add_row(
            cve.cve_id,
            Text(format_cvss(score), style=severity_color(score)),
            f"{cve.epss:.4f}" if cve.epss else "-",
            "[red]YES[/]" if cve.kev else "[dim]no[/]",
            cve.published_time.strftime("%Y-%m-%d") if cve.published_time else "-",
            (cve.summary or "")[:50],
        )

    console.print(table)


def render_cpes(cpes: list[str], fmt: Format):
    if fmt == Format.json:
        print(json.dumps(cpes, indent=2))
        return

    if fmt == Format.tsv:
        for cpe in cpes:
            print(cpe)
        return

    for cpe in cpes:
        console.print(f"[cyan]{cpe}[/]")


@app.command()
def cve(
    cve_id: str = typer.Argument(help="CVE identifier (e.g., CVE-2021-44228)"),
    fmt: Format = typer.Option(None, "--format", "-f", help="Output format"),
):
    """Look up a specific CVE by ID."""
    fmt = get_format(fmt)
    with CVEDBClient() as client:
        try:
            result = client.get_cve(cve_id.upper())
            render_cve_detail(result, fmt)
        except Exception as e:
            err_console.print(f"[red]Error:[/] {e}")
            raise typer.Exit(1)


@app.command()
def cves(
    product: str = typer.Option(None, "--product", "-p", help="Product name to search"),
    cpe: str = typer.Option(None, "--cpe", "-c", help="CPE 2.3 identifier"),
    kev: bool = typer.Option(False, "--kev", "-k", help="Only show CISA KEV entries"),
    sort_epss: bool = typer.Option(False, "--sort-epss", "-e", help="Sort by EPSS score"),
    start: str = typer.Option(None, "--start", "-s", help="Start date (YYYY-MM-DD)"),
    end: str = typer.Option(None, "--end", help="End date (YYYY-MM-DD)"),
    limit: int = typer.Option(25, "--limit", "-l", help="Max results"),
    no_limit: bool = typer.Option(False, "--no-limit", "-L", help="No limit (fetch all)"),
    skip: int = typer.Option(0, "--skip", help="Skip N results"),
    count: bool = typer.Option(False, "--count", help="Only show count"),
    fmt: Format = typer.Option(None, "--format", "-f", help="Output format"),
):
    """Search CVEs by product or CPE."""
    if not product and not cpe:
        err_console.print("[red]Error:[/] Provide --product or --cpe")
        raise typer.Exit(1)

    fmt = get_format(fmt)
    start_date = f"{start}T00:00:00" if start else None
    end_date = f"{end}T23:59:59" if end else None

    with CVEDBClient() as client:
        try:
            result = client.get_cves(
                product=product,
                cpe23=cpe,
                is_kev=kev,
                sort_by_epss=sort_epss,
                start_date=start_date,
                end_date=end_date,
                limit=None if no_limit else limit,
                skip=skip,
                count=count,
            )
            if count:
                if fmt == Format.json:
                    print(json.dumps({"total": result}))
                else:
                    print(result)
            elif result:
                render_cve_table(result, fmt)
            else:
                err_console.print("[yellow]No CVEs found.[/]")
        except Exception as e:
            err_console.print(f"[red]Error:[/] {e}")
            raise typer.Exit(1)


@app.command()
def cpes(
    product: str = typer.Argument(help="Product name to search"),
    limit: int = typer.Option(25, "--limit", "-l", help="Max results"),
    no_limit: bool = typer.Option(False, "--no-limit", "-L", help="No limit (fetch all)"),
    skip: int = typer.Option(0, "--skip", help="Skip N results"),
    count: bool = typer.Option(False, "--count", help="Only show count"),
    fmt: Format = typer.Option(None, "--format", "-f", help="Output format"),
):
    """List CPEs for a product."""
    fmt = get_format(fmt)
    with CVEDBClient() as client:
        try:
            result = client.get_cpes(product=product, limit=None if no_limit else limit, skip=skip, count=count)
            if count:
                if fmt == Format.json:
                    print(json.dumps({"total": result}))
                else:
                    print(result)
            elif result:
                render_cpes(result, fmt)
            else:
                err_console.print("[yellow]No CPEs found.[/]")
        except Exception as e:
            err_console.print(f"[red]Error:[/] {e}")
            raise typer.Exit(1)


if __name__ == "__main__":
    app()