VulnerabilityWeights.py

#
# Copyright (c) 2016-2018, Edgewise Networks Inc. All rights reserved.
#

from elasticsearch import Elasticsearch
from Graph import Graph, Edge

class VulnerabilityWeights(object):

    def __init__(self, defaultWeight=10.0):
        self.defaultWeight = defaultWeight
        try:
            self.es = Elasticsearch([{'host': 'localhost', 'port': 9200}])
        except:
            raise Exception("Start Elastic!")
    
    def mapVulnerabilityScoreToDistance(self, score):
        if score == 0:
            return 100.0
        else:
            return 1.0/score

    def queryVulnerability(self, q):
        sz = self.es.count(index='cve', q=q)
        resp = self.es.search(index='cve', size=min(sz["count"], 10000), q=q)
        maxScore = resp["hits"]["max_score"]
        maxHits = [hit for hit in resp["hits"]["hits"] if hit["_score"] == maxScore]
        keys = {'baseScore', 'exploitabilityScore', 'impactScore'}
        maxVals = {}
        for hit in maxHits:
            doc = hit["_source"]
            for k in keys:
                if k in doc:
                    if k not in maxVals or maxVals[k] < doc[k]:
                        maxVals[k] = doc[k]
        return self.mapVulnerabilityScoreToDistance(maxVals.get('baseScore', 0))

    #def getNmapVulnerability(self, nmapDict):
    #    s = nmapDict["service"]
    #    v = nmapDict["version"]
    #    q = s + " " + v
    #    return self.queryVulnerability(q)

    #def nmapsToGraph(self, nmapList, addVulnWeights=True):
    #    weightedNmaps = self.addWeights(nmapList, addOne=(not addVulnWeights))
    #    g = self.makeGraph(weightedNmaps)
    #    return g

    #def addWeights(self, nmapList, addOne):
    #    if addOne:
    #        for nmap in nmapList:
    #            nmap["weight"] = 1
    #        return nmapList
    #    else:
    #        for nmap in nmapList:
    #            w = self.getNmapVulnerability(nmap)
    #            nmap["weight"] = w
    #        return nmapList

    #def makeGraph(self, nmapList):
    #    es = self.getEdges(nmapList)
    #    g = Graph(es)
    #    return g

    #def getEdges(self, nmapList):
    #    edgeDict = {}
    #    for nmap in nmapList:
    #        src, dst, w = nmap["source_ip"], nmap["dest_ip"], nmap["weight"]
    #        if w is None:
    #            continue
    #        if (src, dst) in edgeDict:
    #            if w < edgeDict[(src, dst)]:
    #                edgeDict[(src, dst)] = w
    #        else:
    #            edgeDict[(src, dst)] = w
    #    return [ Edge(src, dst, w) for (src, dst), w in edgeDict.items() ]