via4 to objects (a20b1448) · Commits · Zhang Yushan / cve-search

lib/Database.py

+25 −1

Original line number	Diff line number	Diff line
		@@ -13,7 +13,7 @@ import pymongo
		import sys

		from lib.Config import Configuration as conf
		from lib.Objects import CVE, CPE, CWE, CAPEC
		from lib.Objects import CVE, CPE, CWE, CAPEC, VIA4
		from lib.Singleton import Singleton

		# Code
		@@ -266,6 +266,20 @@ class Database(metaclass=Singleton):
		return [CAPEC.fromDict(x) for x in
		self.sanitize(self.colCAPEC.find())]

		########
		# VIA4 #
		########
		def via4_upsert(self, data):
		if len(data)>0:
		bulk=self.colVIA4.initialize_unordered_bulk_op()
		for x in data:
		if type(x) is VIA4:
		bulk.find({'id': x.id}).upsert().update({'$set': x.dict()})
		bulk.execute()

		def via4_get(self, cveid):
		via4 = self.colVIA4.find_one({'id': cveid})
		return VIA4.fromDict(via4) if via4 else None

		########
		# Info #
		@@ -321,3 +335,13 @@ class Database(metaclass=Singleton):
		def capec_getUpdate(self):
		return self._getUpdate(self.colCAPEC)

		# VIA4
		def via4_info(self):
		return self._getInfo(self.colVIA4)

		def via4_setUpdate(self, date):
		self._setUpdate(self.colVIA4, date)

		def via4_getUpdate(self):
		return self._getUpdate(self.colVIA4)

+26 −1

Original line number	Diff line number	Diff line
		@@ -16,7 +16,7 @@ from passlib.hash import pbkdf2_sha256

		from lib.Config import Configuration as conf
		from lib.Database import Database
		from lib.Objects import CVE, CPE, CWE, CAPEC
		from lib.Objects import CVE, CPE, CWE, CAPEC, VIA4
		from lib.Singleton import Singleton
		from lib.Toolkit import toStringFormattedCPE

		@@ -31,6 +31,7 @@ class DatabaseLayer(metaclass=Singleton):
		self.CPE = CPEs()
		self.CWE = CWEs()
		self.CAPEC = CAPECs()
		self.VIA4 = VIA4s()
		self.Whitelist = MarkList('whitelist')
		self.Blacklist = MarkList('blacklist')
		self.Users = Users()
		@@ -328,3 +329,27 @@ class CAPECs:
		related_weaknesses.append(rw)
		self.related[w].append(c)
		c.weaknesses = related_weaknesses

		########
		# VIA4 #
		########
		class VIA4s:
		def __init__(self):
		self.db = Database()

		def upsert(self, via4):
		if not isinstance(via4, list):
		via4 = [via4]
		if not all(isinstance(x, VIA4) for x in via4):
		raise ValueError()
		self.db.via4_upsert(via4)

		def get(self, cveid):
		via4 = self.db.via4_get(cveid)
		return via4 if via4 else None

		def updated(self, date=None):
		if date:
		self.db.via4_setUpdate(date)
		else:
		return self.db.via4_getUpdate()

+30 −1

Original line number	Diff line number	Diff line
		@@ -217,3 +217,32 @@ class CAPEC:
		return cls(data['id'], data['name'], data['summary'],
		data['prerequisites'], data['solutions'],
		data['related_weakness'])

		########
		# VIA4 #
		########
		class VIA4:
		def __init__(self, **kwargs):
		for key, val in kwargs.items():
		if not key.startswith("__"):
		setattr(self, key, self._getChild(val))

		def _getChild(self, data):
		if isinstance(data, list):
		values = []
		for val in data:
		values.append(self._getChild(val))
		return values
		return VIA4(**data) if isinstance(data, dict) else data

		def dict(self):
		data = {}
		for key, val in self.__dict__.items():
		if key.startswith("__"): continue
		if isinstance(val, VIA4): val = val.dict()
		data[key] = val
		return data

		@classmethod
		def fromDict(cls, data):
		return cls(**data)