#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Minimal class to get the last entries from the CVE database.
#
# Ranking and CPE lookup are optional.
#
# Software is free software released under the "GNU Affero General Public License v3.0"
#
# Copyright (c) 2012-2018  Alexandre Dulaunoy - a@foo.be
# Copyright (c) 2015-2018  Pieter-Jan Moreels - pieterjan.moreels@gmail.com

import itertools
import os
import sys

from lib.DatabaseLayer import getCVE, sanitize
from lib.cpe_conversion import split_cpe_name

runPath = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(runPath, ".."))


class CveHandler(object):
    def __init__(
        self,
        collection="cves",
        rankinglookup=False,
        namelookup=False,
        capeclookup=False,
        via4lookup=False,
        subscorelookup=False,
    ):
        self.collectionname = collection
        self.rankinglookup = rankinglookup
        self.namelookup = namelookup
        self.capeclookup = capeclookup
        self.subscorelookup = subscorelookup
        self.via4lookup = via4lookup

        self.collection = collection

    def getcapec(self, cweid=None):
        from lib.DatabaseLayer import (
            getCAPECFor,
        )

        if cweid is None or not self.capeclookup:
            return False
        e = getCAPECFor(cweid)
        capec = []
        for f in e:
            capec.append(f)
        return capec

    def getcpe(self, cpeid=None):
        from lib.DatabaseLayer import (
            getCPE,
            getAlternativeCPE,
        )

        if not self.namelookup:
            return cpeid
        e = getCPE(cpeid)
        if e is None:
            e = getAlternativeCPE(cpeid)
            if e is None:
                return cpeid
        if "id" in e:
            return e["title"]

    def getVIA4(self, cveid=None):
        from lib.DatabaseLayer import (
            getVIA4,
        )

        if not self.via4lookup:
            return cveid
        e = getVIA4(cveid)
        return e if e else cveid

    def getCveById(self, cveid):
        if cveid is None:
            return None

        return self.getcve(cveid)

    # This method can be used to standardize a mongo doc
    # after getting it from mongo
    def getCveFromMongoDbDoc(self, cve_doc):
        if cve_doc is None:
            return None

        cve_doc = sanitize(cve_doc)
        if "cwe" in cve_doc and self.capeclookup:
            cwe_list = [
                cwe.split("-")[1] for cwe in cve_doc["cwe"] if cwe.lower() != "unknown"
            ]
            if cwe_list:
                cve_doc["capec"] = self.getcapec(cweid=cwe_list)
        if "vulnerable_configuration" in cve_doc:
            vulconf = []
            ranking = []
            for conf in cve_doc["vulnerable_configuration"]:
                vulconf.append({"id": conf, "title": self.getcpe(cpeid=conf)})
                if self.rankinglookup:
                    rank = self.getranking(cpeid=conf)
                    if rank and rank not in ranking:
                        ranking.append(rank)
            cve_doc["vulnerable_configuration"] = vulconf
        # if self.rankinglookup and len(ranking) > 0:
        #     e["ranking"] = ranking
        if self.via4lookup:
            f = self.getVIA4(cve_doc.get("id"))
            if isinstance(f, dict):
                cve_doc = dict(itertools.chain(cve_doc.items(), f.items()))

        return cve_doc

    def getcve(self, cveid=None):
        return (
            None
            if cveid is None
            else self.getCveFromMongoDbDoc(getCVE(cveid, collection=self.collection))
        )

    def getranking(self, cpeid=None, loosy=True):
        from lib.DatabaseLayer import (
            findRanking,
        )

        if cpeid is None:
            return False

        result = False

        i = None

        if loosy:
            for x in split_cpe_name(cpeid):
                if x != "":
                    i = findRanking(x, regex=True)
                if i is None:
                    continue
                if "rank" in i:
                    result = i["rank"]
        else:
            i = findRanking(cpeid, regex=True)
            if i is None:
                return result
            if "rank" in i:
                result = i["rank"]

        return result

    def get(self, limit=5, skip=0):
        from lib.DatabaseLayer import (
            getCVEs,
        )

        entries = []
        for item in getCVEs(limit=limit, skip=skip, collection=self.collection)[
            "results"
        ]:
            if not self.namelookup and not self.rankinglookup:
                entries.append(item)
            elif self.namelookup or self.rankinglookup:
                if "vulnerable_configuration" in item:
                    vulconf = []
                    ranking = []
                    for conf in item["vulnerable_configuration"]:
                        vulconf.append(self.getcpe(cpeid=conf))
                        if self.rankinglookup:
                            rank = self.getranking(cpeid=conf)
                            if rank and rank not in ranking:
                                ranking.append(rank)
                    item["vulnerable_configuration"] = vulconf
                    if self.rankinglookup:
                        item["ranking"] = ranking
                if "ranking" in item:
                    if len(item["ranking"]) == 0:
                        del item["ranking"]
                if "cwe" in item and self.capeclookup:
                    capec_ids = []
                    for cwe in item["cwe"]:
                        if cwe.lower() != "unknown":
                            try:
                                capec_ids.extend(self.getcapec(cweid=cwe.split("-")[1]))
                            except (IndexError, ValueError):
                                continue
                    if capec_ids:
                        item["capec"] = capec_ids

                entries.append(item)

        return entries
