cve_search.py

#!/usr/bin/env python
import argparse
import json
import re
import sys
import zipfile
from datetime import datetime
from functools import partial
from pathlib import Path

import requests

REPO_URL = "https://github.com/CVEProject/cvelistV5/archive/refs/heads/main.zip"
ZIP_FILE_PATH = Path(__file__).parent / "main.zip"
CVE_JSON_RE = re.compile(r".+/CVE-\d{4}-\d{4,}\.json")

CSI = "\x1b["
RED = CSI + "31m"
GREEN = CSI + "32m"
YELLOW = CSI + "33m"
BLUE = CSI + "34m"
RESET = CSI + "0m"

print_err = partial(print, file=sys.stderr)


def download_repo() -> None:
    print_err(f"{YELLOW}Downloading repository from {REPO_URL}...{RESET}")
    response = requests.get(
        REPO_URL,
        headers={
            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
        },
    )
    response.raise_for_status()
    with open(ZIP_FILE_PATH, "wb") as f:
        f.write(response.content)
    print_err(f"{GREEN}ZIP archive saved to {ZIP_FILE_PATH}{RESET}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Search for CVE files in the CVEProject repository."
    )
    parser.add_argument(
        "-s",
        "--search",
        required=True,
        type=str,
        help="Search term for the CVE description (fuzzy search supported).",
    )
    parser.add_argument(
        "-m",
        "--min-score",
        type=float,
        default=9.5,
        help="Minimum CVSS score to filter CVE files.",
    )
    parser.add_argument(
        "-a",
        "--date-after",
        type=str,
        help="Filter CVE by datePublished after (format: YYYY-MM-DD).",
    )
    parser.add_argument(
        "-b",
        "--date-before",
        type=str,
        help="Filter CVE by datePublished before (format: YYYY-MM-DD).",
    )
    args = parser.parse_args()
    search = args.search.lower()

    # Парсинг дат
    date_after = (
        datetime.strptime(args.date_after, "%Y-%m-%d")
        if args.date_after
        else None
    )
    date_before = (
        datetime.strptime(args.date_before, "%Y-%m-%d")
        if args.date_before
        else None
    )

    if not ZIP_FILE_PATH.exists():
        download_repo()

    with zipfile.ZipFile(ZIP_FILE_PATH, "r") as zip:
        for file_info in zip.infolist():
            if not CVE_JSON_RE.fullmatch(file_info.filename):
                continue
            with zip.open(file_info) as f:
                try:
                    data = json.load(f)
                    cna = data.get("containers", {}).get("cna", {})
                    descriptions = cna.get("descriptions", [])
                    metrics = cna.get("metrics", [])
                    base_score = get_base_score(metrics)

                    # Проверка даты публикации
                    date_published = data["cveMetadata"].get("datePublished")
                    if date_published:
                        published_date = datetime.strptime(
                            date_published.split("T")[0], "%Y-%m-%d"
                        )

                        # Фильтр date-after
                        if date_after and published_date < date_after:
                            continue

                        # Фильтр date-before
                        if date_before and published_date > date_before:
                            continue

                    if (
                        any(
                            search in desc.get("value", "").lower()
                            for desc in descriptions
                        )
                        and base_score >= args.min_score
                    ):
                        cve_id = data["cveMetadata"]["cveId"]
                        description = next(
                            (
                                desc["value"]
                                for desc in descriptions
                                if desc.get("lang") == "en"
                            ),
                            "No description available",
                        )
                        print(cve_id)
                        print(f"Score: {base_score}")
                        print(date_published)
                        print()
                        print(description)
                        print()
                except (json.JSONDecodeError, KeyError):
                    print_err(
                        f"{RED}Error reading file: {file_info.filename}{RESET}"
                    )


def get_base_score(metrics: list[dict]) -> float:
    for metric in metrics:
        match metric:
            case {"cvssV3_1": cvss}:
                return cvss.get("baseScore", 0.0)
            case {"cvssV3_0": cvss}:
                return cvss.get("baseScore", 0.0)
    return 0.0


if __name__ == "__main__":
    main()