# SPDX-License-Identifier: LGPL-2.0-or-later
# SPDX-FileCopyrightText: 2024 Ics inc.
# SPDX-FileContributor: Sergey Missan <smissan@ics.com>
# SPDX-FileContributor: Michael Dingwall <mdingwall@ics.com>
# SPDX-FileContributor: Chris Rizzitello <crizzitello@ics.com>
# SPDX-FileContributor: Qin Zhang <qzhang@ics.com>
# SPDX-FileContributor: Boris Ralchenko <bralchenko@ics.com>

import sqlite3
import time
import json
import os
import requests
import argparse

from typing import Final
from pathlib import Path
from datetime import datetime
from cpeparser import CpeParser
from tqdm import tqdm

from ics_sbom_libs.common.vulnerability import Vulnerability

# Setup Logging
import logging

log = logging.getLogger(__name__)

FETCH_INTERVAL_SEC: Final = 2
FETCH_INTERVAL_SEC_NO_API_KEY: Final = 6
CVE_RECORD_PER_PAGE: Final = 2000
CPE_RECORD_PER_PAGE: Final = 10000
# CVE v2. See https://nvd.nist.gov/developers/vulnerabilities for refrence
NIST_BASE_URL: Final = "https://services.nvd.nist.gov/rest/json/{}/2.0/"
CVE_URL: Final = NIST_BASE_URL.format("cves")
CPE_URL: Final = NIST_BASE_URL.format("cpes")

_cache_dir = Path("~").expanduser() / ".cache" / "icsbom"

# see cve_schema.sql for the version numbers.
_nvd_db_version = "4.0"
_nvd_dbFile = f"nvd_v{_nvd_db_version}.db"
_api_key = ""  # this will for it to load the saved api_key if one is saved.


class VulnerabilityDatabase:
    _cache_dir: Path
    _db_file_name: str
    _api_key: str | None

    def __init__(self, cache_dir=_cache_dir, db_file=_nvd_dbFile, api_key=_api_key):
        self.cache_dir = cache_dir
        self.db_file_name = db_file
        self.api_key = api_key

        self._last_query: datetime | None = None

        self._setup()

    def __del__(self):
        if self.con:
            self.con.close()

    @property
    def cache_dir(self):
        return self._cache_dir

    @cache_dir.setter
    def cache_dir(self, new_dir):
        if new_dir is None:
            return

        if isinstance(new_dir, str):
            self._cache_dir = Path(new_dir)
        elif isinstance(new_dir, Path):
            self._cache_dir = new_dir

        # Handle the possibility that '~' was used in the path.
        self._cache_dir = self._cache_dir.expanduser()

        # In the event there are any symbolic links in the path, let's just resolve everything out.
        self._cache_dir = self._cache_dir.resolve()

    @property
    def db_file_name(self):
        return self._db_file_name

    @db_file_name.setter
    def db_file_name(self, new_file):
        if isinstance(new_file, str):
            self._db_file_name = new_file

    @property
    def api_key(self):
        return self._api_key

    @api_key.setter
    def api_key(self, new_key):
        if not new_key:
            self._api_key = self._load_api_key()
            return

        if isinstance(new_key, str):
            if new_key == "none":
                self._api_key = ""
            else:
                self._api_key = new_key

    def _load_api_key(self):
        if not self.cache_dir or not self.cache_dir.exists():
            return ""

        api_key_path = self.cache_dir / "api_key.txt"
        if not api_key_path.exists():
            return ""

        api_key = api_key_path.read_text().strip()
        if not api_key:
            return ""

        return api_key

    def _save_api_key(self):
        if not self.cache_dir or not self.cache_dir.exists() or not self.api_key:
            return

        api_key_path = self.cache_dir / "api_key.txt"
        api_key_path.write_text(f"{self.api_key}")

    @property
    def db_path(self):
        return self.cache_dir / self.db_file_name

    def _setup(self):
        if not self.cache_dir.exists():
            self.cache_dir.mkdir(parents=True)

        if not os.path.exists(self.db_path):
            self._init_db_()

        self.con = sqlite3.connect(self.db_path)

    @staticmethod
    def setup_args(parser: argparse.ArgumentParser):
        if not parser:
            return

        parser.add_argument("--cache_dir", type=str, default=_cache_dir, help="Cache Directory: contains the database")
        parser.add_argument("--api_key", type=str, default=_api_key, required=False, help="API key for NVD database")
        parser.add_argument(
            "--save_key",
            action="store_true",
            required=False,
            help="Save the API key for the NVD database to a cache file in cache_dir.",
        )
        parser.add_argument(
            "--db_file",
            type=str,
            default=f"{_nvd_dbFile}",
            help="Database file path. If file does not exist, a new database will be generated "
            "from scratch and saved to the given path. This will take sometime.",
        )

    def process_args(self, args):
        if not args:
            return

        self.cache_dir = args.cache_dir
        self.db_file_name = args.db_file
        self.api_key = args.api_key

        if args.save_key and args.api_key != "none":
            self._save_api_key()
            args.save_key = False

        self._setup()

    def _init_db_(self):
        with open(Path(__file__).parent.resolve() / "cve_schema.sql") as fp:
            self.con = sqlite3.connect(self.db_path)
            cur = self.con.cursor()
            cur.executescript(fp.read())

    def _get_latest_timestamp_(self):
        try:
            curs = self.con.execute("SELECT last_modified FROM cve_severity ORDER BY last_modified DESC LIMIT 1;")
            latest = curs.fetchone()[0]
            curs.close()
            return latest
        except Exception:
            return None

    def _get_status_value(self, key: str, default=None):
        try:
            curs = self.con.execute("SELECT value FROM status WHERE key = ?", (key,))
            item: str = curs.fetchone()
            curs.close()
            if item is None:
                return default
            return item[0]
        except Exception:
            return default

    def _set_status_value(self, key: str, value):
        try:
            curs = self.con.cursor()
            curs.execute("REPLACE INTO status (key, value) VALUES (?,?)", (key, f"{value}"))
            self.con.commit()
        except Exception:
            return

    def _remove_status_key(self, key: str):
        try:
            curs = self.con.cursor()
            curs.execute("DELETE FROM status WHERE key = ?", (key,))
            self.con.commit()
        except Exception:
            return

    def _query_nvd(self, nvd_url: str, query: str):
        query_url = f"{nvd_url}?"

        query_interval = FETCH_INTERVAL_SEC if self.api_key else FETCH_INTERVAL_SEC_NO_API_KEY
        query_delta = (datetime.now() - self._last_query).seconds if self._last_query else query_interval + 1
        if query_delta < query_interval:
            time.sleep(query_interval - query_delta)

        if self.api_key:
            response = requests.get(query_url + query, headers={"apiKey": self.api_key})
        else:
            response = requests.get(query_url + query)

        self._last_query = datetime.now()

        return response

    def query_cve_from_nvd(self, cve_id: str):
        if not cve_id:
            return None

        cve_query = f"cveId={cve_id}"
        retry_count = 0
        finished = False
        data = None
        while True:

            try:
                response = self._query_nvd(CVE_URL, cve_query)
            except Exception as e:
                log.exception(f"{e}")
                break

            if response.status_code == 503:
                time.sleep(6)
                retry_count += 1
                continue
            else:
                retry_count = 0

            try:
                data = response.text
                finished = True

            except Exception as e:
                log.exception(str(e) + "\n" + "\nstatus_code:" + str(response.status_code))
                data = None
                break

            if finished:
                break

        return data

    def create_database(self):

        def get_data(src_name: str, src_url, data_handler, records_per_page):
            # Grab a minimal amount of data just to get the totals.
            download_db = not bool(self._get_status_value(f"initial_{src_name}_download_completed", False))

            start_index = int(self._get_status_value(f"last_{src_name}_record_received", 0)) if download_db else 0
            last_update = self._get_status_value(f"{src_name}_last_updated")
            last_update_query = ""
            if not download_db and last_update is not None:
                last_update = last_update.replace("'", "")
                diff = datetime.utcnow() - datetime.fromisoformat(last_update)
                if diff.days > 100:
                    self._init_db_()  # Not allow to have long duration data fetch
                else:
                    log.info(f"Updating {src_name.upper()} records since last update: {last_update}")
                    last_update_query = f"lastModStartDate={last_update}&lastModEndDate=" + str(
                        datetime.utcnow().isoformat()
                    )

            pbar = tqdm(total=1, desc=f"Updating NVD {src_name.upper()} Database", unit_scale=1, unit="Records")
            pbar.update(start_index)
            retry_count = 0
            finished = False
            while True:

                try:
                    response = self._query_nvd(
                        src_url, last_update_query + f"&resultsPerPage={records_per_page}&startIndex={start_index}"
                    )
                except Exception as e:
                    log.exception(f"{e}")
                    break

                if response.status_code == 503:
                    pbar.set_postfix_str("Request Failed...")
                    time.sleep(6)
                    retry_count += 1
                    pbar.set_postfix_str(f"Request Failed... Retrying (Retry Count: {retry_count})")
                    continue
                else:
                    pbar.set_postfix_str(refresh=False)
                    retry_count = 0

                try:
                    data = json.loads(response.text)
                    data_handler(data)
                    if data["startIndex"] + data["resultsPerPage"] >= data["totalResults"]:
                        finished = True

                    else:
                        start_index = data["startIndex"] + data["resultsPerPage"]

                except Exception as e:
                    log.exception(str(e) + "\n" + "\nstatus_code:" + str(response.status_code))
                    break

                if pbar.total != data["totalResults"]:
                    pbar.total = data["totalResults"]
                    pbar.refresh()

                # self.con.commit()

                if not finished:
                    self._set_status_value(f"last_{src_name}_record_received", start_index)
                else:
                    self._remove_status_key(f"last_{src_name}_record_received")
                    self._set_status_value(f"initial_{src_name}_download_completed", finished)

                self._set_status_value(f"{src_name}_last_updated", datetime.utcnow().isoformat())

                self.con.commit()
                pbar.update(data["resultsPerPage"])

                if finished:
                    break

            pbar.close()

        get_data("cve", CVE_URL, self._process_cve_data_, CVE_RECORD_PER_PAGE)
        get_data("cpe", CPE_URL, self._process_cpe_data_, CPE_RECORD_PER_PAGE)

        self.con.close()

    def _process_cve_data_(self, data):
        cveArray = []
        rangeArray = []
        weaknessArray = []
        cves = data["vulnerabilities"]
        for cve in cves:
            ranges = []
            weakness = []
            config_block = []
            data = CveDataHelper.get_data(cve["cve"], ranges, weakness, config_block)

            if data is not None:
                cveArray.append(data)
                if len(ranges) > 0:
                    rangeArray.extend(ranges)
                if len(weakness) > 0:
                    weaknessArray.extend(weakness)

        curs = self.con.cursor()
        vuln_names = Vulnerability.sql_query_name_list()
        curs.executemany(
            "INSERT INTO cve_severity (" + ", ".join(vuln_names) + ") "
            "VALUES (" + ", ".join(f":{name}" for name in vuln_names[0:-1]) + ", 'NVD') "
            "ON CONFLICT (cve_number) DO UPDATE SET " + ", ".join(f"{name}=:{name}" for name in vuln_names[1:-1]) + ";",
            cveArray,
        )

        # Clear ranges and weakness for CVE records
        cveNums = list()
        for item in cveArray:
            cveNums.append("'" + item["cve_number"] + "'")
        cveNums = ",".join(cveNums)
        curs.execute(f"DELETE FROM cve_range WHERE cve_number IN ({cveNums})")
        curs.execute(f"DELETE FROM cve_weakness WHERE cve_number IN ({cveNums})")

        # TODO, limit duplicated ranges.
        if len(rangeArray) > 0:
            key_list = [
                "cve_number",
                "vendor",
                "product",
                "version",
                "cpe",
                "vulnerable",
                "versionStartIncluding",
                "versionStartExcluding",
                "versionEndIncluding",
                "versionEndExcluding",
                "part_type",
                "data_source",
            ]
            sql_str = (
                "INSERT INTO cve_range ("
                + ", ".join(key_list)
                + ") "
                + " VALUES ("
                + ", ".join(f":{name}" for name in key_list[0:-2])
                + ", :part , 'NVD') "
            )
            curs.executemany(sql_str, rangeArray)

        if len(weaknessArray) > 0:
            curs.executemany(
                "INSERT INTO cve_weakness (cve_number, value) VALUES (:cve_number, :value) "
                "ON CONFLICT (cve_number, value) DO NOTHING",
                weaknessArray,
            )

    def _process_cpe_data_(self, data):
        cpeArray = []
        products = data["products"]
        for product in products:
            # Are there multiple products in the CWE NVD table?
            if "cpe" in product.keys():
                cpeArray.append(CpeRecordDataHelper.parse_data(product["cpe"]))

        if not cpeArray:
            return

        curs = self.con.cursor()
        cpe_properties = list(cpeArray[0].keys())
        curs.executemany(
            "INSERT INTO cpe_dictionary ("
            + ", ".join(cpe_properties)
            + ") "
            + "VALUES ("
            + ", ".join(f":{name}" for name in cpe_properties)
            + ") "
            + "ON CONFLICT (cpe_id) DO UPDATE SET "
            + ", ".join(f"{name}=:{name}" for name in cpe_properties[1:])
            + ";",
            cpeArray,
        )

    def query_cache(self, query: str, parameters=()):
        if not query:
            return None

        return self.con.execute(query, parameters)

    def query_cpe_dictionary(self, package_name: str):
        query = f"SELECT cpe FROM cpe_dictionary WHERE product='{package_name}'"
        cpe_strings = []

        cursor = self.con.execute(query)
        if not cursor:
            return cpe_strings

        res = cursor.fetchall()
        if not res:
            return cpe_strings

        for cpe in res:
            cpe_strings.append(cpe[0])

        return cpe_strings

    def get_cve(self, cve_id: str):
        if not self.con:
            return None

        results = self.query_cache(
            "SELECT " + ", ".join(Vulnerability.sql_query_name_list()) + " FROM cve_severity WHERE cve_number=?",
            (cve_id,),
        )
        cve = Vulnerability(results.fetchone())

        results = self.query_cache("SELECT value FROM cve_weakness WHERE cve_number=?", (cve_id,))
        cve.cwes = [cwe[0] for cwe in results.fetchall()]

        return cve


class CveDataHelper:
    range_fields: Final = [
        "versionStartIncluding",
        "versionStartExcluding",
        "versionEndIncluding",
        "versionEndExcluding",
    ]

    @staticmethod
    def get_data(cve, ranges, weakness, config_block):
        desc = cve["descriptions"][0]["value"]
        if desc.startswith("** REJECT **") or desc.startswith("Rejected reason:"):
            return None
        cve_number = cve["id"]
        dateTime = cve["lastModified"] if "lastModified" in cve else cve["published"]
        data = {"cve_number": cve_number, "description": desc, "last_modified": dateTime}
        CveDataHelper.get_metrics(cve, data)
        CveDataHelper.get_range(cve, cve_number, ranges)
        CveDataHelper.get_configurations(cve, config_block)
        CveDataHelper.get_weakness(cve, cve_number, weakness)
        data["configurations"] = json.dumps(config_block)

        return data

    @staticmethod
    def parse_node(cve, cve_number) -> []:
        ret = []
        if "cpeMatch" not in cve:
            return ret
        cpe = CpeParser()

        for cpeMatch in cve["cpeMatch"]:
            result = cpe.parser(cpeMatch["criteria"])
            result["cve_number"] = cve_number
            result["cpe"] = cpeMatch["criteria"]
            result["vulnerable"] = cpeMatch["vulnerable"]
            for field in CveDataHelper.range_fields:
                if field in cpeMatch:
                    result[field] = cpeMatch[field]
                else:
                    result[field] = ""
            ret.append(result)
        return ret

    @staticmethod
    def get_weakness(cve, cve_number, weakness) -> None:
        if "weaknesses" not in cve:
            return

        for item in cve["weaknesses"]:
            weakness.append({"value": item["description"][0]["value"], "cve_number": cve_number})

    @staticmethod
    def get_range(cve, cve_number, ranges):
        if "configurations" not in cve:
            return

        for item in cve["configurations"]:  # TODO Check the configurations array. And the operator in the array.
            for node in item["nodes"]:
                ranges.extend(CveDataHelper.parse_node(node, cve_number))

    @staticmethod
    def get_configurations(cve, config_block):
        if "configurations" not in cve:
            return

        config_block.extend(cve["configurations"])

    @staticmethod
    def get_metrics(cve, data):
        if "cvssMetricV31" in cve["metrics"]:
            data["severity"] = cve["metrics"]["cvssMetricV31"][0]["cvssData"]["baseSeverity"]
            data["score"] = cve["metrics"]["cvssMetricV31"][0]["cvssData"]["baseScore"]
            data["cvss_vector"] = cve["metrics"]["cvssMetricV31"][0]["cvssData"]["vectorString"]
            data["cvss_version"] = float(cve["metrics"]["cvssMetricV31"][0]["cvssData"]["version"])
        elif "cvssMetricV30" in cve["metrics"]:
            data["severity"] = cve["metrics"]["cvssMetricV30"][0]["cvssData"]["baseSeverity"]
            data["score"] = cve["metrics"]["cvssMetricV30"][0]["cvssData"]["baseScore"]
            data["cvss_vector"] = cve["metrics"]["cvssMetricV30"][0]["cvssData"]["vectorString"]
            data["cvss_version"] = float(cve["metrics"]["cvssMetricV30"][0]["cvssData"]["version"])
        elif "cvssMetricV2" in cve["metrics"]:
            data["severity"] = cve["metrics"]["cvssMetricV2"][0]["baseSeverity"]
            data["score"] = cve["metrics"]["cvssMetricV2"][0]["cvssData"]["baseScore"]
            data["cvss_vector"] = cve["metrics"]["cvssMetricV2"][0]["cvssData"]["vectorString"]
            data["cvss_version"] = float(cve["metrics"]["cvssMetricV2"][0]["cvssData"]["version"])
        else:
            data["severity"] = "unknown"
            data["score"] = None
            data["cvss_vector"] = "unknown"
            data["cvss_version"] = 0


class CpeRecordDataHelper:
    property_map: Final = {
        "cpe_id": "cpeNameId",
        "cpe": "cpeName",
        "deprecated": "deprecated",
        "created": "created",
        "last_modified": "lastModified",
    }

    @staticmethod
    def parse_data(data):
        cpe_record = CpeRecordDataHelper.parse_cpe(data[CpeRecordDataHelper.property_map["cpe"]])
        for prop in CpeRecordDataHelper.property_map:
            cpe_record[prop] = data[CpeRecordDataHelper.property_map[prop]]

        return cpe_record

    @staticmethod
    def parse_cpe(cpe_name):
        result = CpeParser().parser(cpe_name)
        return {"vendor": result["vendor"], "product": result["product"]}
