cli.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
from datetime import datetime, timezone
import json
import logging
import os
import shutil
import sys
import types

from rich.console import Console
from rich.live import Live
from rich.table import Table

from vdb.lib import config, display, db6 as db_lib, search
from vdb.lib.aqua import AquaSource
from vdb.lib.cve_model import CVE
from vdb.lib.gha import GitHubSource
from vdb.lib.osv import OSVSource

if sys.platform == "win32" and os.environ.get('PYTHONIOENCODING') is None:
    sys.stdin.reconfigure(encoding="utf-8")
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")

ORAS_AVAILABLE = False
# oras is an optional dependency
try:
    from vdb.lib.orasclient import download_image

    ORAS_AVAILABLE = True
except ImportError:
    pass

console = Console()

logging.basicConfig(
    level=logging.INFO, format="%(levelname)s [%(asctime)s] %(message)s"
)
LOG = logging.getLogger(__name__)
for _ in ("httpx",):
    logging.getLogger(_).disabled = True

AT_LOGO = r"""
             ___
  /\  ._  ._  | |_  ._ _   _. _|_
 /--\ |_) |_) | | | | (/_ (_|  |_
      |   |
"""


def build_args():
    """
    Constructs command line arguments for the vulnerability-db tool
    """
    parser = argparse.ArgumentParser(
        description="AppThreat's vulnerability database and package search library with a sqlite storage."
    )
    parser.add_argument(
        "--clean",
        action="store_true",
        default=False,
        dest="clean",
        help="Clear the vulnerability database cache from platform specific user_data_dir.",
    )
    parser.add_argument(
        "--cache",
        action="store_true",
        default=False,
        dest="cache",
        help="Cache vulnerability information in platform specific user_data_dir.",
    )
    parser.add_argument(
        "--cache-os",
        action="store_true",
        default=False,
        dest="cache_os",
        help="Cache OS vulnerability information in platform specific user_data_dir.",
    )
    parser.add_argument(
        "--only-osv",
        action="store_true",
        default=False,
        dest="only_osv",
        help="Use only OSV as the source. Use with --cache.",
    )
    parser.add_argument(
        "--only-aqua",
        action="store_true",
        default=False,
        dest="only_aqua",
        help="Use only Aqua vuln-list as the source. Use with --cache.",
    )
    parser.add_argument(
        "--only-ghsa",
        action="store_true",
        default=False,
        dest="only_ghsa",
        help="Use only recent ghsa as the source. Use with --cache.",
    )
    parser.add_argument(
        "--search",
        dest="search",
        help="Search for the package or vulnerability ID (CVE, GHSA, ALSA, DSA, etc.) in the database. Use purl, cpe, or git http url.",
    )
    parser.add_argument(
        "--list-malware",
        action="store_true",
        default=False,
        dest="list_malware",
        help="List latest malwares with CVE ID beginning with MAL-.",
    )
    parser.add_argument(
        "--bom",
        dest="bom_file",
        help="Search for packages in the CycloneDX BOM file.",
    )
    parser.add_argument(
        "--download-image",
        action="store_true",
        default=False,
        dest="download_image",
        help="Downloaded pre-created vdb image to platform specific user_data_dir. Application vulnerabilities only.",
    )
    parser.add_argument(
        "--download-full-image",
        action="store_true",
        default=False,
        dest="download_full_image",
        help="Downloaded pre-created vdb image to platform specific user_data_dir. All vulnerabilities including OS.",
    )
    parser.add_argument(
        "--print-vdb-metadata",
        action="store_true",
        default=False,
        dest="print_vdb_metadata",
        help="Display metadata about the current vdb in user_data_dir.",
    )
    return parser.parse_args()


def print_results(results):
    added_row_keys = {}
    table = Table(title="VDB Results", show_lines=True)
    table.add_column("CVE", justify="left", max_width=20)
    table.add_column("Locator")
    table.add_column("Description")
    table.add_column("Affected Symbols", max_width=50)
    if isinstance(results, types.GeneratorType):
        with Live(
            table, console=console, refresh_per_second=4, vertical_overflow="visible"
        ):
            for result_gen in results:
                if isinstance(result_gen, dict):
                    display.add_table_row(table, result_gen, added_row_keys)
                elif isinstance(result_gen, list):
                    for res in result_gen:
                        display.add_table_row(table, res, added_row_keys)
                elif isinstance(result_gen, types.GeneratorType):
                    for res in result_gen:
                        display.add_table_row(table, res, added_row_keys)
    elif isinstance(results, list):
        for res in results:
            display.add_table_row(table, res, added_row_keys)
        console.print(table)


def create_db_file_metadata(sources, cve_data_count, cve_index_count):
    """Method to create the vdb file metadata"""
    metadata = {
        "created_utc": datetime.now(tz=timezone.utc).isoformat(
            timespec="seconds"
        ),
        "cve_data_count": cve_data_count,
        "cve_index_count": cve_index_count,
        "sources": [s.__class__.__name__ for s in sources],
    }
    include_list = []
    ignore_list = []
    # Collect injected metadata from environment variables
    for name, value in os.environ.items():
        if value is None:
            continue
        if name.startswith("VDB_METADATA_"):
            if value in ("true", "1"):
                value = True
            elif value in ("false", "0"):
                value = False
            metadata[name.replace("VDB_METADATA_", "").lower()] = value
        elif name.startswith("VDB_IGNORE_") and value in ("true", "1"):
            ignore_list.append(name.replace("VDB_IGNORE_", "").lower())
        elif name.startswith("VDB_EXCLUDE_") and value in ("true", "1"):
            ignore_list.append(name.replace("VDB_EXCLUDE_", "").lower())
        elif name.startswith("VDB_INCLUDE_") and value in ("true", "1"):
            include_list.append(name.replace("VDB_INCLUDE_", "").lower())
    if include_list:
        metadata["include_list"] = include_list
    if ignore_list:
        metadata["ignore_list"] = ignore_list
    if os.getenv("NVD_START_YEAR"):
        metadata["start_year"] = os.getenv("NVD_START_YEAR")
    if os.getenv("GITHUB_PAGE_COUNT"):
        metadata["github_page_count"] = os.getenv("GITHUB_PAGE_COUNT")
    return metadata


def print_db_file_metadata(metadata_file):
    if not os.path.exists(metadata_file):
        return
    with open(metadata_file, encoding="utf-8") as fp:
        db_meta = json.load(fp)
        table = Table(title="VDB Summary", show_lines=True, caption=f"Metadata file: {metadata_file}")
        table.add_column("Property")
        table.add_column("Value")
        for k, v in db_meta.items():
            table.add_row(k, str(v))
        console.print(table)


def main():
    """Main function"""
    args = build_args()
    print(AT_LOGO)
    if args.clean:
        if os.path.exists(config.DATA_DIR):
            shutil.rmtree(config.DATA_DIR, ignore_errors=True)
        if not os.path.exists(config.DATA_DIR):
            LOG.info("VDB cache cleaned successfully.")
        else:
            LOG.info("VDB cache at %s not cleaned successfully.", config.DATA_DIR)
    if args.print_vdb_metadata:
        print_db_file_metadata(config.VDB_METADATA_FILE)
    if args.download_image or args.download_full_image:
        db_url = config.VDB_DATABASE_URL if args.download_full_image else config.VDB_APP_ONLY_DATABASE_URL
        if ORAS_AVAILABLE:
            LOG.info(
                "Downloading vdb image from %s to %s",
                db_url,
                config.DATA_DIR,
            )
            try:
                download_image(db_url, config.DATA_DIR)
                print_db_file_metadata(config.VDB_METADATA_FILE)
            except Exception as e:
                LOG.info("Vulnerability database download using oras failed. Please try again.")
                LOG.exception(e)
                return
        else:
            console.print(
                "Oras library is not available. Install using 'pip install appthreat-vulnerability-db[oras]' and re-run this command.",
                markup=False
            )
    elif args.cache or args.cache_os:
        db_lib.get()
        db_lib.clear_all()
        if args.only_osv:
            sources = [OSVSource()]
        elif args.only_aqua:
            sources = [AquaSource()]
        elif args.only_ghsa:
            sources = [GitHubSource()]
        else:
            sources = [OSVSource(), GitHubSource()]
        # AquaSource also includes NVD
        if args.cache_os and not args.only_aqua:
            sources.insert(0, AquaSource())
        for s in sources:
            LOG.info("Refreshing %s", s.__class__.__name__)
            s.refresh()
        cve_data_count, cve_index_count = db_lib.stats()
        console.print(
            "cve_data_count", cve_data_count, "cve_index_count", cve_index_count
        )
        db_lib.optimize_and_close_all()
        # Create metadata about the database
        with open(config.VDB_METADATA_FILE, mode="w", encoding="utf-8") as meta_file:
            json.dump(
                create_db_file_metadata(sources, cve_data_count, cve_index_count),
                meta_file,
            )
    if args.search:
        if db_lib.needs_update():
            console.print(
                "Vulnerability database needs to be refreshed. Please run 'vdb --download-image' to download the latest."
            )
        results = search.search_by_any(args.search, with_data=True)
        if results:
            print_results(results)
        else:
            console.print("No results found!")
    elif args.bom_file:
        if os.path.exists(args.bom_file):
            results_generator = search.search_by_cdx_bom(args.bom_file, with_data=True)
            print_results(results_generator)
    elif args.list_malware:
        results_generator = search.latest_malware(with_data=True)
        print_results(results_generator)


if __name__ == "__main__":
    main()