more database layer abstraction (f81b8914) · Commits · Zhang Yushan / cve-search

bin/search_xmpp.py

+0 −3

Original line number	Diff line number	Diff line
		@@ -46,9 +46,6 @@ runPath = os.path.dirname(os.path.realpath(__file__))

		rankinglookup = True

		db = Configuration.getMongoConnection()
		collection = db.cves

		helpmessage = "\nlast <n> cve entries (output: JSON) \n"
		helpmessage = helpmessage + "cvetweet <n> cve entries (output: Text) \n"
		helpmessage = helpmessage + "browse vendors and products (output: JSON)\n"

lib/DatabaseLayer.py

+17 −1

Original line number	Diff line number	Diff line
		@@ -22,7 +22,7 @@ colWHITELIST=db['mgmt_whitelist']
		colBLACKLIST=db['mgmt_blacklist']
		colUSERS= db['mgmt_users']
		colSEEN= db['mgmt_seen']

		colINFO= db['info']
		# Functions
		def sanitize(x):
		if type(x)==pymongo.cursor.Cursor:
		@@ -33,6 +33,16 @@ def sanitize(x):
		return x

		# DB Functions
		def setColUpdate(collection, date):
		colINFO.update({"db": collection}, {"$set": {"last-modified": date}}, upsert=True)

		def insertCVE(cve):
		colCVE.insert(cve)

		def updateCVE(cve):
		colCVE.update({"id": cve['id']}, {"$set": {"cvss": cve['cvss'], "summary": cve['summary'], "references": cve[' references'],
		"cwe": cve['cwe'], "vulnerable_configuration": cve['vulnerable_configuration'],
		"vulnerable_configuration_cpe_2_2": cve['vulnerable_configuration_cpe_2_2'], 'last-modified': cve['Modified']}})

		# API Functions
		def cvesForCPE(cpe):
		@@ -81,6 +91,9 @@ def getCVEs(limit=-1, query=[], skip=0):
		cves=colCVE.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)
		return sanitize(cves)

		def getCVE(id):
		return sanitize(colCVE.find_one({"id": id}))

		def getCPE(id):
		return sanitize(colCPE.find_one({"id": id}))

		@@ -90,6 +103,9 @@ def getAlternativeCPE(id):
		def getFreeText(text):
		return [x["obj"] for x in db.command("text", "cves", search=text)["results"]]

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

		# Dynamic data
		def getWhitelist():
		return sanitize(colWHITELIST.find())

sbin/db_mgmt.py

+11 −14

Original line number	Diff line number	Diff line
		@@ -22,6 +22,7 @@ from xml.sax.handler import ContentHandler
		from lib.ProgressBar import progressbar
		from lib.Toolkit import toStringFormattedCPE
		from lib.Config import Configuration
		import lib.DatabaseLayer as dbLayer

		# parse command line arguments
		argparser = argparse.ArgumentParser(description='populate/update the local CVE database')
		@@ -200,12 +201,12 @@ if __name__ == '__main__':
		f = Configuration.getFile(Configuration.getCVEDict() + getfile)
		except:
		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(Configuration.getCVEDict() + getfile))
		i = info.find_one({'db': 'cve'})
		i = dbLayer.getInfo("cve")
		if i is not None:
		if f.headers['last-modified'] == i['last-modified']:
		print("Not modified")
		sys.exit(0)
		info.update({'db': 'cve'}, {"$set": {'last-modified': f.headers['last-modified']}}, upsert=True)
		dbLayer.setColUpdate("cve", f.headers['last-modified'])

		# get your parser on !!
		parser = make_parser()
		@@ -214,18 +215,16 @@ if __name__ == '__main__':
		parser.parse(f)
		for item in ch.cves:
		# check if the CVE already exists.
		x = collection.find({'id': item['id']})
		x = dbLayer.getCVE(item['id'])
		# if so, update the entry.
		if x.count() > 0:
		if x:
		if 'cvss' not in item:
		item['cvss'] = defaultvalue['cvss']
		if 'cwe' not in item:
		item['cwe'] = defaultvalue['cwe']
		collection.update({'id': item['id']}, {'$set': {'cvss': item['cvss'], 'summary': item['summary'], 'references': item['references'],
		'cwe': item['cwe'], 'vulnerable_configuration': item['vulnerable_configuration'],
		'vulnerable_configuration_cpe_2_2':item['vulnerable_configuration_cpe_2_2'], 'last-modified': item['Modified']}})
		dbLayerr.updateCVE(item)
		else:
		collection.insert(item)
		dbLayer.insertCVE(item)
		# get the 'recent' file
		getfile = file_prefix + file_rec + file_suffix
		try:
		@@ -238,9 +237,9 @@ if __name__ == '__main__':
		parser.parse(f)
		for item in progressbar(ch.cves):
		# check if the CVE already exists.
		x = collection.find({'id': item['id']})
		x = dbeLayer.getCVE(item['id'])
		# if so, update the entry.
		if x.count() > 0:
		if x:
		if args.v:
		print("item found : " + item['id'])
		if 'cvss' not in item:
		@@ -249,12 +248,10 @@ if __name__ == '__main__':
		item['cvss'] = float(item['cvss'])
		if 'cwe' not in item:
		item['cwe'] = defaultvalue['cwe']
		collection.update({'id': item['id']}, {'$set': {'cvss': item['cvss'], 'summary': item['summary'], 'references': item['references'],
		'cwe': item['cwe'], 'vulnerable_configuration': item['vulnerable_configuration'],
		'vulnerable_configuration_cpe_2_2':item['vulnerable_configuration_cpe_2_2'], 'last-modified': item['Modified']}})
		dbLayer.updateCVE(item)
		# if not, create it.
		else:
		collection.insert(item)
		dbLayer.insertCVE(item)
		elif args.p:
		# populate is pretty straight-forward, just grab all the files from NVD
		# and dump them into a DB.

Admin message