vulnerability.py

from enum import Enum
from typing import Dict, List


class MatchReason(Enum):
    PRODUCT_MATCH = (8, "product_match")
    VERSION_IN_RANGE = (7, "version_in_range")
    GENERAL_PRODUCT_OK = (6, "general_product_ok")
    DESCRIPTION_MATCH = (5, "description_match")
    GENERAL_PRODUCT_UNCERTAIN = (4, "general_product_uncertain")
    SINGLE_HIGHER_VERSION = (3, "single_higher_version")
    VULN_ID = (2, "vuln_id")
    N_A = (1, "n_a")

    def __init__(self, priority, description):
        self._priority = priority
        self._description = description

    def __lt__(self, other):
        if isinstance(other, MatchReason):
            return self._priority < other._priority
        return NotImplemented

    def __eq__(self, other):
        if isinstance(other, MatchReason):
            return self._priority == other._priority
        return NotImplemented

    def __str__(self):
        return self._description

    def __repr__(self):
        return str(self)


class Vulnerability:
    def __init__(
        self,
        _id: str,
        match_reason: MatchReason,
        match_sources: List[str],
        description="",
        published="",
        modified="",
        cvss_ver="",
        cvss="",
        cvss_vec="",
        exploits: List[str] = None,
        cisa_known_exploited=False,
        href: str = "",
        aliases: Dict[str, str] = None,
        tracked_by=None,
        epss="",
        reported_patched_by=None,
        **misc,
    ):
        """
        Create a new vulnerability.

        Args:
            _id (str): Unique identifier for this vulnerability.
            match_reason (MatchReason): Reason for why this vulnerability was matched.
            match_sources (List[str]): List of database sources contributing to the match.
            description (str, optional): Description of this vulnerability. Defaults to an empty string.
            published (str, optional): Date when this vulnerability was published. Defaults to an empty string.
            modified (str, optional): Date when this vulnerability was last modified. Defaults to an empty string.
            cvss_ver (str, optional): CVSS version used for scoring. Defaults to an empty string.
            cvss (str, optional): CVSS base score or string representation. Defaults to an empty string.
            cvss_vec (str, optional): CVSS vector string. Defaults to an empty string.
            exploits (List[str], optional): List containing URLs to known exploits. Defaults to None.
            cisa_known_exploited (bool, optional): Whether this vulnerability is known to be exploited per CISA. Defaults to False.
            href (str, optional): A hyperlink reference to this vulnerability; will be stored under aliases. Defaults to an empty string.
            aliases (Dict[str, str], optional): Alternate IDs for this vulnerability, including references if available. Defaults to None.
            tracked_by (optional): List of data sources tracking this vulnerability. Defaults to None.
            epss (str, optional): EPSS score or string representation. Defaults to an empty string.
            reported_patched_by (List[str]): List of database sources reporting the vulnerability patched for the queried version.
            **misc: Any additional keyword arguments for extended metadata.
        """

        if not _id:
            raise ValueError("Vuln ID cannot be None")
        if not match_reason:
            raise ValueError("Vuln match_reason cannot be None")
        if not match_sources:
            raise ValueError("Vuln match_sources cannot be None")

        self.id = _id
        self.match_reason = match_reason
        if isinstance(match_sources, list):
            self.match_sources = match_sources.copy()
        else:
            self.match_sources = [match_sources]
        self.description = description
        self.published = published
        self.modified = modified
        self.cvss_ver = cvss_ver
        self.cvss = cvss
        self.cvss_vec = cvss_vec
        self.exploits = set()
        if exploits:
            self.add_exploits(exploits)
        self.cisa_known_exploited = cisa_known_exploited
        if isinstance(aliases, list):
            self.aliases = {}
            for alias in aliases:
                self.aliases[alias] = ""
        else:
            self.aliases = dict(aliases) if aliases else {}
        if _id not in self.aliases or not self.aliases[_id]:
            self.aliases[_id] = href
        self.tracked_by = list(tracked_by) if tracked_by else []
        self.epss = epss
        if reported_patched_by is None:
            self.reported_patched_by = []
        elif isinstance(reported_patched_by, list):
            self.reported_patched_by = reported_patched_by.copy()
        else:
            self.reported_patched_by = [reported_patched_by]
        self.misc = dict(misc) if misc else {}

    def to_dict(self):
        def convert(value):
            if isinstance(value, (str, int, float, bool, type(None), list, dict)):
                return value
            if isinstance(value, set):
                return list(value)
            return str(value)  # Stringify everything else (e.g., custom objects)

        return {k: convert(v) for k, v in self.__dict__.items()}

    def __str__(self):
        return str(self.to_dict())

    def __repr__(self):
        return str(self)

    def merge_with_vulnerability(self, other):
        if not isinstance(other, self.__class__):
            raise TypeError("Can only merge with the same class type")

        # copy over simple values of other vuln
        for key, value in other.__dict__.items():
            if key in (
                "match_reason",
                "match_sources",
                "exploits",
                "aliases",
                "tracked_by",
                "misc",
                "cvss_ver",
                "cvss",
                "cvss_vec",
                "epss",
                "reported_patched_by",
            ):
                continue
            # always copy attribute if this vuln doesn't have it
            if value and not getattr(self, key):
                setattr(self, key, value)
            # copy attribute if it's valid and other has a better match reason
            if other.match_reason > self.match_reason and value:
                setattr(self, key, value)

        # only copy other's CVSS if self doesn't have one or other's match_reason is higher
        copy_cvss = False
        if not self.is_cvss_valid() and other.is_cvss_valid():
            copy_cvss = True
        elif other.match_reason > self.match_reason and other.is_cvss_valid():
            copy_cvss = True
        if copy_cvss:
            self.cvss_ver = other.cvss_ver
            self.cvss = other.cvss
            self.cvss_vec = other.cvss_vec

        # only copy other's EPSS if self doesn't have one or other's match_reason is higher
        if not self.is_epss_valid() and other.is_epss_valid():
            self.epss = other.epss
        elif other.match_reason > self.match_reason and other.is_epss_valid():
            self.epss = other.epss

        # increase match reason if data source of other vuln is more certain
        if other.match_reason > self.match_reason:
            self.match_reason = other.match_reason

        # merge complicated values
        self.match_sources = list(set(self.match_sources + other.match_sources))
        self.tracked_by = list(set(self.tracked_by + other.tracked_by))
        self.add_exploits(other.exploits)
        self.reported_patched_by = list(
            set(self.reported_patched_by + other.reported_patched_by)
        )

        # merge aliases (ID of a vuln will always have be stored as alias already)
        new_aliases = {}
        for alias, href in self.aliases.items():
            new_aliases[alias] = href
        for alias, href in other.aliases.items():
            if alias not in new_aliases or not new_aliases[alias]:
                new_aliases[alias] = href
        self.aliases = new_aliases

        # merge misc items up to one level, otherwise discard
        for key, value in other.misc.items():
            if key not in self.misc:
                self.misc[key] = value
            elif isinstance(self.misc[key], list) and isinstance(value, list):
                self.misc[key] = list(set(self.misc[key] + value))
            elif isinstance(self.misc[key], dict) and isinstance(value, dict):
                for key2, val2 in value.items():
                    if key2 not in self.misc[key]:
                        self.misc[key][key2] = val2

    def is_cvss_valid(self):
        if not self.cvss_ver:
            return False
        if not self.cvss:
            return False
        if not self.cvss_vec:
            return False

        try:
            cvss_ver = float(self.cvss_ver)
            cvss = float(self.cvss)
            if cvss_ver < 0:
                return False
            if cvss < 0:
                return False
        except ValueError:
            return False

        return True

    def is_epss_valid(self):
        if not self.epss:
            return False
        try:
            epss = float(self.epss)
            if epss < 0 or epss > 1:
                return False
        except:
            return False
        return True

    def add_alias(self, alias, href="", source_db=""):
        if alias not in self.aliases or not self.aliases[href]:
            self.aliases[alias] = href
        if source_db:
            self.add_tracked_by(source_db)

    def add_tracked_by(self, source_db):
        if source_db not in self.tracked_by:
            self.tracked_by.append(source_db)

    def add_exploits(self, exploits):
        if (not isinstance(exploits, list)) and (not isinstance(exploits, set)):
            exploits = [exploits]
        for exploit in exploits:
            if exploit.endswith("/"):
                if exploit[:-1] in self.exploits:
                    continue
                elif exploit.startswith("http:") and "https:" + exploit[5:] in self.exploits:
                    continue
                if exploit + ".git" not in self.exploits:
                    self.exploits.add(exploit)
            else:
                if exploit.startswith("http:") and "https:" + exploit[5:] in self.exploits:
                    continue
                if exploit + "/" in self.exploits:
                    self.exploits.remove(exploit + "/")
                if exploit + ".git" not in self.exploits:
                    self.exploits.add(exploit)

    # Getters and setters
    def get_vuln_id(self):
        return self.vuln_id

    def set_vuln_id(self, value):
        self.vuln_id = value

    def get_match_reason(self):
        return self.match_reason

    def set_match_reason(self, value):
        self.match_reason = value

    def get_match_sources(self):
        return self.match_sources

    def set_match_sources(self, value):
        self.match_sources = value

    def get_description(self):
        return self.description

    def set_description(self, value):
        self.description = value

    def get_published(self):
        return self.published

    def set_published(self, value):
        self.published = value

    def get_modified(self):
        return self.modified

    def set_modified(self, value):
        self.modified = value

    def get_href(self):
        return self.href

    def set_href(self, value):
        self.href = value

    def get_cvss_ver(self):
        return self.cvss_ver

    def set_cvss_ver(self, value):
        self.cvss_ver = value

    def get_cvss(self):
        return self.cvss

    def set_cvss(self, value):
        self.cvss = value

    def get_cvss_vec(self):
        return self.cvss_vec

    def set_cvss_vec(self, value):
        self.cvss_vec = value

    def get_cisa_known_exploited(self):
        return self.cisa_known_exploited

    def set_cisa_known_exploited(self, value):
        self.cisa_known_exploited = value

    def get_aliases(self):
        return self.aliases

    def set_aliases(self, value):
        self.aliases = value

    def get_tracked_by(self):
        return self.tracked_by

    def set_tracked_by(self, value):
        self.tracked_by = value

    def get_epss(self):
        return self.epss

    def set_epss(self, value):
        self.epss = value

    def set_patched(self, source):
        if source not in self.reported_patched_by:
            self.reported_patched_by.append(source)

    def is_patched(self):
        return bool(self.reported_patched_by)