vulnerability_node.py

from bs4 import BeautifulSoup
from pymongo import MongoClient

class VulnerabilityNode(object):
    def __init__(self, CVEName, port, entry=False):
        self.vulnerabilityName = CVEName
        self.vulnerabilityPort = port
        self.mapping1 = { 'High': 2, 'Low': 1, 'None': 0 }
        self.mapping2 = { 'Admin': 2, 'User': 1, 'None': 0 }
        self.type = 'vuln'
        vulns = MongoClient('localhost', 27017).project.vulnerabilities
        query = {}
        query['cveName'] = CVEName
        result = list(vulns.find(query))[0]
        self.accessVector = result['access_vector']
        self.accessLevel = result['gained_access']
        self.requiredPrivilege = result['required_priv']
        self.entry = entry

    # def get_gained_access_and_access_vector(self):
    #     url = "https://www.cvedetails.com/cve/" + self.vulnerabilityName
    #     req = Request(url, headers={ 'User-Agent': 'Mozilla/5.0' })
    #     html_doc = urlopen(req).read()
    #     soup = BeautifulSoup(html_doc, 'lxml')
    #     table = soup.find("table", { 'id': 'cvssscorestable', 'class': 'details' })
    #     field_row = table.findAll("tr")[6]
    #     field_value = field_row.find("span").string
    #     return self.mapping2[field_value]
    
    # def get_required_privilege(self):
    #     url = "https://nvd.nist.gov/vuln/detail/" + self.vulnerabilityName
    #     html_doc = urlopen(url)
    #     soup = BeautifulSoup(html_doc, 'lxml')
    #     tag = soup.find('span', { 'data-testid': 'vuln-cvssv3-pr' })
    #     if tag:
    #         field_value = tag.string.strip()
    #     else:
    #         field_value = "None" # By default, "None" privileges are required
    #     return self.mapping1[field_value]

    def to_string(self):
        return "({}, {}, {}, {}, {})".format(self.vulnerabilityName, self.vulnerabilityPort, self.requiredPrivilege, self.accessLevel, self.type)

    def __eq__(self, other):
        return self.vulnerabilityName == other.vulnerabilityName and self.vulnerabilityPort == other.vulnerabilityPort

    def __hash__(self):
        return hash(('name', self.vulnerabilityName, 'port', self.vulnerabilityPort))

# v = VulnerabilityNode('CVE-2012-3137', 53)
# print(v.accessLevel)
# print(v.accessVector)
# print(v.requiredPrivilege)