<reponame>Nmap-Vulnerability-Enumeration/nve-console-client
import json
import utils

class Vulnerability:
    def __init__(self,
                 cve_id: str,
                 published: str = None,
                 patched: bool = None,          # patch status
                 configurations: list = None,
                 impact=None
                 ):
        self.cve_id = cve_id
        self.published_date = published
        self.patched = patched
        self.configs = configurations
        self.impact = impact

    @staticmethod
    def from_nist(entry):
        if len(entry) == 0:
            raise ValueError("_dict cannot be empty")

        cve_id = entry["cve"]["CVE_data_meta"]["ID"]
        published = entry["publishedDate"] if "publishedDate" in entry else None
        description = None
        impact = entry["impact"] if "impact" in entry else None

        # get description
        for description_data in entry["cve"]["description"]["description_data"]:
            if description_data["lang"] == "en":
                description = description_data["value"]
                break

        configs = entry["configurations"] if "configurations" in entry else None

        return Vulnerability(
            cve_id,
            published=published,
            patched=None,
            configurations=configs,
            impact=impact
        )

    def applicable_to(self, cpe_list):
        if self.configs == None or "nodes" not in self.configs:
            return False

        return utils.is_vulnerable(cpe_list, self.configs["nodes"])

    def is_vulnerable(self, device):
        if self.configs == None or "nodes" not in self.configs:
            return False

        cpe_list = device.get_all_cpes()

        for node in self.configs["nodes"]:
            if self._is_vuln_helper(node, cpe_list):
                return True
        return False

    def _is_vuln_helper(self, config_node, cpe_list):
        if len(cpe_list) == 0:
            return False

        and_node = config_node["operator"] == "AND" if "operator" in config_node else False
        negate = config_node["negate"] if "negate" in config_node else False
        return_val = None

        if "cpe_match" in config_node and len(config_node["cpe_match"]) > 0:
            for match in config_node["cpe_match"]:
                cpe_found = utils.cpe_in_list(match["cpe23Uri"], cpe_list)
                if cpe_found and not and_node:  # if cpe found and its an or node
                    return_val = True
                    break
                elif (not cpe_found) and and_node:  # if cpe not found and its an and node
                    return_val = False
                    break
            if return_val == None:
                return_val = and_node  # if its an and node, return true; otherwise return false

        elif "children" in config_node and len(config_node["children"]) > 0:
            for child in config_node["children"]:
                is_vuln = self._is_vuln_helper(child, cpe_list)
                if is_vuln and not and_node:
                    return_val = True
                    break
                elif (not is_vuln) and and_node:
                    return_val = False
                    break
            if return_val == None:
                return_val = and_node

        else:
            raise ValueError("Node needs children or cpe_match")

        return return_val if negate == False else (not return_val)

    def enum_configs(self):
        if self.configs == None or "nodes" not in self.configs:
            return []

        configs = []
        for node in self.configs["nodes"]:
            child_enum = self._enum_configs_helper(node)
            for enum in child_enum:
                if len(enum) > 0:
                    configs.append(enum)
        
        return configs

    def _enum_configs_helper(self, current_node):
        _configs = []
        and_node = current_node["operator"] == "AND" if "operator" in current_node else False
        negate = current_node["negate"] if "negate" in current_node else False

        if "cpe_match" in current_node and len(current_node["cpe_match"]) > 0:
            _temp = []
            for match in current_node["cpe_match"]:
                if and_node:
                    _temp.append(match["cpe23Uri"])
                else:
                    _configs.append([match["cpe23Uri"]])
            if and_node:
                _configs.append(_temp)

        elif "children" in current_node and len(current_node["children"]) > 0:
            for child in current_node["children"]:
                child_enum = self._enum_configs_helper(child)
                if len(child_enum) == 0:
                    continue

                elif and_node: # and case
                    '''
                    merge every item already found with every item in child_num
                    '''
                    if len(_configs) == 0:
                        _configs = child_enum
                    else:
                        _temp = []
                        for child_item in child_enum:
                            if len(child_item) > 0:
                                for item in _configs:                    
                                    _temp.append(item + child_item)
                        _configs = _temp
                else: # Or case
                    '''
                    Append new found items to configs
                    '''
                    _configs = _configs + child_enum

        else:
            raise ValueError("Node needs children or cpe_match")

        return _configs


class VulnerabilityEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, Vulnerability):
            return {
                "_type": "vuln",
                "value": {
                    "cve": o.cve_id,
                    "published": o.published_date,
                    "patched": o.patched,
                    "configs": o.configs,
                    "impact": o.impact
                }
            }

        else:
            return super().default(o)
