Commit a6a3a0fe authored by PidgeyL's avatar PidgeyL
Browse files

more objects

parent 133e6207

lib/Database.py

+40 −13
Original line number Diff line number Diff line
@@ -81,6 +81,34 @@ class Database(metaclass=Singleton): 
    return x





  ############

  # Database #

  ############

  def db_getStats(self, include_admin=False):

    def getSize(collection):

      return self.db[collection].count()



    def getInfo(collection):

      return self.sanitize(self.colINFO.find_one({"db": collection}))



    def getLastModified(collection):

      info=getInfo(collection)

      return info['last-modified'] if info else None





    data={'cves': {}, 'cpe': {}, 'cpeOther': {}, 'capec': {}, 'cwe': {}, 'via4': {}}

    for key in data.keys():

      data[key] = {'size': getSize(key.lower()),

                   'last_update': getLastModified(key.lower())}

    if include_admin:

      data['whitelist']={'size': self.colWHITELIST.count()}

      data['blacklist']={'size': self.colBLACKLIST.count()}

      data = {'stats': {'size_on_disk': self.db.command("dbstats")['storageSize'],

                        'db_size':      self.db.command('dbstats')['dataSize'],

                        'name':         conf.getDatabaseName()},

              'data':  data}

    return data



  #######################

  # Blacklist/Whitelist #

  #######################
@@ -207,18 +235,12 @@ class Database(metaclass=Singleton): 
    cves = self.colCVE.find(query).sort(sort[0], sort[1]).limit(limit).skip(skip)

    return [CVE.fromDict(x) for x in self.sanitize(cves)]



  def getCVEs(limit=False, query=[], skip=0, cves=None, collection=None):

    col=colCVE if not collection else db[collection]

    if type(query) == dict: query=[query]

    if type(cves) == list: query.append({"id": {"$in": cves}})

    if len(query) == 0:

      cve=col.find().sort("Modified", -1).limit(limit).skip(skip)

    elif len(query)  == 1:

      cve=col.find(query[0]).sort("Modified", -1).limit(limit).skip(skip)

    else:

      cve=col.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)

    return sanitize(cve) or []



  def cve_textSearch(self, text):

    try: # Before Mongo 3

      data = [x["obj"] for x in self.db.command("text", "cves", search=text)["results"]]

    except: # As of Mongo 3

      data = self.sanitize(self.colCVE.find({"$text":{"$search":text}}))

    return [CVE.fromDict(x) for x in data]



  ########

  # CPEs #
@@ -285,9 +307,14 @@ class Database(metaclass=Singleton): 
      bulk.execute()



  def via4_get(self, cveid):

    via4 = self.colVIA4.find_one({'id': cveid})

    via4 = self.sanitize(self.colVIA4.find_one({'id': cveid}))

    via4.pop("id")

    return VIA4.fromDict(via4) if via4 else None



  def via4_link(self, key, val):

    cveList=[x['id'] for x in colVIA4.find({key: val})]

    return self.sanitize(cve_query(query={'id':{'$in':cveList}}))



  ########

  # Info #

  ########

lib/DatabaseLayer2.py

+9 −0
Original line number Diff line number Diff line
@@ -36,6 +36,9 @@ class DatabaseLayer(metaclass=Singleton): 
    self.Blacklist = MarkList('blacklist')

    self.Users     = Users()



  def db_info(self, include_admin=False):

    return self.db.db_getStats(include_admin)



####################

# Black-/Whitelist #

####################
@@ -239,6 +242,9 @@ class CVEs: 
  def forCPE(self, cpe):

    return self.db.cve_forCPE(cpe)



  def textSearch(self, search):

    return self.db.cve_textSearch(search)



  # Info

  def info(self):

    return self.db.getCVEInfo()
@@ -379,3 +385,6 @@ class VIA4s: 
      self.db.via4_setUpdate(date)

    else:

      return self.db.via4_getUpdate()



  def link(self, key, value):

    return self.db.via4_link(key, value)

lib/Objects.py

+19 −18
Original line number Diff line number Diff line
@@ -76,7 +76,7 @@ class CVE: 
                 published, modified=None, impact=None, access=None,

                 cwe=None, references=[], cvss_time=None):

        tk.assertType(str, id=id, summary=summary)

        tk.assertType(float, cvss=cvss)

        tk.assertType((float, str, None), cvss=cvss)

        tk.assertType(datetime.datetime, published=published)

        tk.assertType((datetime.datetime, None), modified=modified,

                                                 cvss_time=cvss_time)
@@ -90,7 +90,7 @@ class CVE: 




        self.id                       = id.upper()

        self.cvss                     = cvss

        self.cvss                     = cvss and float(cvss) or None

        self.summary                  = summary

        self.vulnerable_configuration = vulnerable_configuration

        self.published                = published
@@ -109,8 +109,8 @@ class CVE: 
                'vulnerable_configuration': vuln_conf,

                'Published':                self.published,

                'Modified':                 self.modified,

                'impact':                   self.impact.dict(),

                'access':                   self.access.dict(),

                'impact':                   self.impact and self.impact.dict() or None,

                'access':                   self.access and self.access.dict() or None,

                'cwe':                      "CWE-"+self.cwe.id,

                'references':               self.references,

                'cvss-time':                self.cvss_time}
@@ -119,8 +119,8 @@ class CVE: 
    def fromDict(cls, data):

        # Creating dud CPEs with the correct ID, to be replaced

        vc  = [CPE(x) for x in data['vulnerable_configuration']]

        i   = data['impact'] and Impact.fromDict(data['impact']) or None

        a   = data['access'] and Access.fromDict(data['access']) or None

        i   = data.get('impact') and Impact.fromDict(data['impact']) or None

        a   = data.get('access') and Access.fromDict(data['access']) or None

        if data.get('cwe'): # Get ensures backwards compatibility

          # Create a dud cwe with just the ID, so we can assign the

           #  pointer later
@@ -128,7 +128,7 @@ class CVE: 
        else: cwe = None

        return cls(data['id'], data['cvss'], data['summary'], vc,

                   data['Published'], data['Modified'], i, a, cwe,

                   data['references'], data['cvss-time'])

                   data['references'], data.get('cvss-time'))





#######
@@ -138,13 +138,14 @@ class CPE: 
    def __init__(self, id, title=None, references=[]):

        tk.assertType(str, id=id)

        tk.assertType((str, None),  title=title)

        tk.assertType((list, tuple), references=references)

        tk.assertType((list, tuple, None), references=references)

        if references:

            tk.assertTypeForAllIn(str, references)



        self.id         = tk.toStringFormattedCPE(id)

        self.id_2_2     = tk.toOldCPE(id)

        self.title      = title if title else tk.cpeTitle(self.id)

        self.references = list(references)

        self.references = references and list(references) or []



    def dict(self):

        return {'id':         self.id,
@@ -154,7 +155,7 @@ class CPE: 


    @classmethod

    def fromDict(cls, data):

        return cls(data['id'], data['title'], data['references'])

        return cls(data['id'], data['title'], data.get('references'))



#######

# CWE #
@@ -174,14 +175,14 @@ class CWE: 
    def dict(self):

        return {'id':                  self.id,

                'name':                self.name,

                'description': self.description,

                'description_summary': self.description,

                'status':              self.status,

                'weakness':    self.weakness}

                'weaknessabs':         self.weakness}



    @classmethod

    def fromDict(cls, data):

        return cls(data['id'], data['name'], data['description'],

                   data['status'], data['weakness'])

        return cls(data['id'], data['name'], data['description_summary'],

                   data['status'], data['weaknessabs'])



#########

# CAPEC #

CRA Git | Maintained and supported by SUSTech CRA and CCSE