Initial CVE to Object (ca25e397) · Commits · Zhang Yushan / cve-search

lib/Database.py

+7 −6

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

		from lib.Objects import CPE, CWE, CAPEC
		from lib.Objects import CVE, CPE, CWE, CAPEC
		from lib.Singleton import Singleton

		# Code
		@@ -171,17 +171,18 @@ class Database(metaclass=Singleton):
		# CVEs #
		########
		# Data manipulation
		def upsertCVE(self, data):
		if type(data) is dict: data = [data]
		def cve_upsert(self, data):
		if type(data) is CVE: data = [data]
		if len(data)>0:
		bulk=self.colCVE.initialize_unordered_bulk_op()
		for x in data:
		bulk.find({'id': x['id']}).upsert().update({'$set': x})
		bulk.find({'id': x.id}).upsert().update({'$set': x.dict()})
		bulk.execute()

		# Data retrieval
		def getCVE(self, id):
		return sanitize(col.find_one({"id": id}))
		def cve_get(self, id):
		cve = self.sanitize(self.colCVE.find_one({"id": id}))
		return CVE.fromDict(cve) if cve else None

		def cvesForCPE(cpe):
		if not cpe: return []

lib/DatabaseLayer2.py

+12 −3

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 CPE, CWE, CAPEC
		from lib.Objects import CVE, CPE, CWE, CAPEC
		from lib.Singleton import Singleton
		from lib.Toolkit import toStringFormattedCPE

		@@ -27,6 +27,7 @@ class DatabaseLayer(metaclass=Singleton):
		self.salt_size = salt_size

		self.db = _db and Database(db=_db) or Database()
		self.CVE = CVEs()
		self.CPE = CPEs()
		self.CWE = CWEs()
		self.CAPEC = CAPECs()
		@@ -197,11 +198,17 @@ class CVEs:
		cve = [cve]
		if not all(isinstance(x, CVE) for x in cve):
		raise ValueError()
		self.db.upsertCVE(cve)
		self.db.cve_upsert(cve)

		# Data retrieval
		def get(self, cveID):
		return self.db.getCVE(cveID)
		cve = self.db.cve_get(cveID.upper())
		# Replace the dud at cve reconstruction time with the pointer
		if cve:
		if cve.cwe: cve.cwe = DatabaseLayer().CWE.get(cve.cwe.id)
		cve.vulnerable_configuration = [DatabaseLayer().CPE.get(x.id)
		for x in cve.vulnerable_configuration]
		return cve

		def forCPE(self, cpe):
		return self.db.cvesForCPE(cpe)
		@@ -277,6 +284,7 @@ class CWEs:
		self.cwe = {x.id: x for x in self.db.cwe_getAll()}
		for c in self.cwe.values():
		c.capec = DatabaseLayer().CAPEC.relatedTo(c.id)
		print(self.cwe)


		#########
		@@ -321,3 +329,4 @@ class CAPECs:
		related_weaknesses.append(rw)
		self.related[w].append(c)
		c.weaknesses = related_weaknesses
		print(self.related)

lib/Objects.py

+29 −5

Original line number	Diff line number	Diff line
		@@ -38,6 +38,12 @@ class Impact:
		'integrity': self.integrity,
		'availability': self.availability}

		@classmethod
		def fromDict(cls, data):
		return cls(data['confidentiality'], data['integrity'],
		data['availability'])


		class Access:
		def __init__(self, complexity, authentication, vector):
		tk.assertType(str, complexity=complexity,
		@@ -59,6 +65,12 @@ class Access:
		'authentication': self.authentication,
		'vector': self.vector}

		@classmethod
		def fromDict(cls, data):
		return cls(data['complexity'], data['authentication'],
		data['vector'])


		class CVE:
		def __init__(self, id, cvss, summary, vulnerable_configuration,
		published, modified=None, impact=None, access=None,
		@@ -70,9 +82,9 @@ class CVE:
		cvss_time=cvss_time)
		tk.assertType(list, vulnerable_configuration=vulnerable_configuration,
		references=references)
		tk.assertType(Impact, impact=impact)
		tk.assertType(Access, access=access)
		tk.assertType(CWE, cwe=cwe)
		tk.assertType((Impact, None), impact=impact)
		tk.assertType((Access, None), access=access)
		tk.assertType((CWE, None), cwe=cwe)
		tk.assertTypeForAllIn(CPE, vulnerable_configuration)
		tk.assertTypeForAllIn(str, references)

		@@ -105,7 +117,18 @@ class CVE:

		@classmethod
		def fromDict(cls, data):
		return cls(data['id'], data['title'], data['references'])
		# 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
		if data.get('cwe'): # Get ensures backwards compatibility
		# Create a dud cwe with just the ID, so we can assign the
		# pointer later
		cwe = CWE(data['cwe'].strip("CWE-"), 'dud', 'dud', 'dud', 'dud')
		else: cwe = None
		return cls(data['id'], data['cvss'], data['summary'], vc,
		data['Published'], data['Modified'], i, a, cwe,
		data['references'], data['cvss-time'])


		#######
		@@ -113,7 +136,8 @@ class CVE:
		#######
		class CPE:
		def __init__(self, id, title=None, references=[]):
		tk.assertType(str, id=id, title=title)
		tk.assertType(str, id=id)
		tk.assertType((str, None), title=title)
		tk.assertType((list, tuple), references=references)
		tk.assertTypeForAllIn(str, references)

test/test_DatabaseLayer.py

+66 −5

Original line number	Diff line number	Diff line
		@@ -106,7 +106,8 @@ class TestDatabase(unittest.TestCase):
		self.capec1 = CAPEC(id="10000", name="test_capec",
		summary="no summary",
		prerequisites="No prerequisites",
		solutions="There's no solution", weaknesses=[])
		solutions="There's no solution",
		weaknesses=["10000"])
		self.cwe1 = CWE(id="10000", name="test_cwe",
		description="test cwe", status="testing",
		weakness='Testing')
		@@ -135,17 +136,77 @@ class TestDatabase(unittest.TestCase):
		self.db.CAPEC.upsert(self.capec1)
		result = self.db.CAPEC.get(self.capec1.id)
		assert self.capec1.id == result.id
		assert self.capec1.summary == result.summary
		assert self.capec1.solutions == result.solutions

		def test_insert_cwe(self):
		self.db.CWE.upsert(self.cwe1)
		result = self.db.CWE.get(self.cwe1.id)
		assert self.cwe1.id == result.id

		assert self.cwe1.name == result.name
		assert self.cwe1.description == result.description

		def test_insert_cpe(self):
		self.db.CPE.upsert([self.cpe1, self.cpe2])
		result = self.db.CPE.get(self.cpe2.id)
		assert self.cpe2.id == result.id
		assert self.cpe2.title == result.title


		def test_insert_cve(self):
		self.db.CVE.upsert(self.cve)
		result = self.db.CVE.get(self.cve1.id)
		assert self.cve1.id == result.id
		assert self.cve1.summary == result.summary
		assert self.cve1.cvss == result.cvss

		class TestDatabase(unittest.TestCase):
		# Unittest functions
		def setUp(self):
		self.db = DatabaseLayer(_db="cvedb_test")
		self.capec1 = CAPEC(id="10000", name="test_capec",
		summary="no summary",
		prerequisites="No prerequisites",
		solutions="There's no solution",
		weaknesses=["10000"])
		self.cwe1 = CWE(id="10000", name="test_cwe",
		description="test cwe", status="testing",
		weakness='Testing')
		self.cpe1 = CPE(id="cpe:/a:test:test1", title="Test CPE 1",
		references=[])
		self.cpe2 = CPE(id="cpe:2.3:a:test:test2", title="Test CPE 2",
		references=[])
		self.cve1 = CVE(id="CVE-0001-0001", cvss=0.1,
		summary="Test Vulnerability",
		vulnerable_configuration=[self.cpe1, self.cpe2],
		published=datetime.datetime(2017,1,1),
		impact=Impact("None", "None", "None"),
		access=Access("Low", "None", "Local"),
		cwe=self.cwe1)

		def tearDown(self):
		del self.db
		del self.capec1
		del self.cwe1
		del self.cpe1
		del self.cpe2
		del self.cve1


		def test_reconstruction_cve(self):
		self.db.CAPEC.upsert(self.capec1)
		self.db.CWE.upsert(self.cwe1)
		self.db.CPE.upsert([self.cpe1, self.cpe2])
		self.db.CVE.upsert(self.cve1)

		result = self.db.CVE.get(self.cve1.id)

		assert self.cve1.cwe.name == result.cwe.name
		assert self.capec1.name == result.cwe.capec[0].name
		# Needed because we can't compare object IDs
		vc1 = [(x.id, x.title) for x in
		self.cve1.vulnerable_configuration]
		vc2 = [(x.id, x.title) for x in result.vulnerable_configuration]
		self.assertListEqual(vc1, vc2)
		assert self.cve1.impact.integrity == result.impact.integrity
		assert self.cve1.access.vector == result.access.vector

Admin message