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

lib/CVEs.py

+3 −7

Original line number	Original line	Diff line number	Diff line
	@@ -41,17 +41,13 @@ class last:

	if rankinglookup:		if rankinglookup:
	self.ranking = connectdb['ranking']		self.ranking = connectdb['ranking']
	if vfeedlookup:
	self.vfeed = connectdb['vfeed']
	if capeclookup:
	self.capec = connectdb['capec']
	if reflookup:		if reflookup:
	self.ref = Configuration.getRedisRefConnection()		self.ref = Configuration.getRedisRefConnection()

	def getcapec(self, cweid=None):		def getcapec(self, cweid=None):
	if cweid is None or not self.capeclookup:		if cweid is None or not self.capeclookup:
	return False		return False
	e = self.capec.find({'related_weakness': {'$in': [cweid, cweid]}})		e = db.getCAPECFor(cweid)
	capec = []		capec = []
	for f in e:		for f in e:
	f.pop('_id')		f.pop('_id')
	@@ -77,10 +73,10 @@ class last:

	def getvfeed(self, cveid=None):		def getvfeed(self, cveid=None):

	if not(self.vfeed):		if not(self.vfeedlookup):
	return cveid		return cveid

	e = self.vfeed.find_one({'id': cveid})		e = db.getvFeed(cveid)

	if e is None:		if e is None:
	return cveid		return cveid

lib/DatabaseLayer.py

+15 −10

Original line number	Original line	Diff line number	Diff line
	@@ -63,13 +63,6 @@ def bulkUpdate(collection, data):
	bulk.find({'id': x['id']}).upsert().update({'$set': x})		bulk.find({'id': x['id']}).upsert().update({'$set': x})
	bulk.execute()		bulk.execute()

	def cweBulkUpdate(cwelist):
	bulk=colCPE.initialize_ordered_bulk_op()
	for x in cwelist:
	# To check: Can we not just update({'$set': x})
	bulk.find({'id': x['id']}).upsert().update({"$set": {'name': x['name'], 'id': x['id'], 'status': x['status'], 'weaknessabs': x['weaknessabs']}})
	bulk.execute()

	def cpeotherBulkInsert(cpeotherlist):		def cpeotherBulkInsert(cpeotherlist):
	colCPEOTHER.insert(cpeotherlist)		colCPEOTHER.insert(cpeotherlist)

	@@ -180,17 +173,29 @@ def getCVE(id):
	def getCPE(id):		def getCPE(id):
	return sanitize(colCPE.find_one({"id": id}))		return sanitize(colCPE.find_one({"id": id}))

			def getCPEs():
			return sanitize(colCPE.find())

			def getAlternativeCPE(id):
			return sanitize(colCPEOTHER.find_one({"id": id}))

			def getAlternativeCPEs():
			return sanitize(colCPEOTHER.find())

			def getvFeed(id):
			return sanitize(colVFEED.find_one({'id': cveid}))

	def getCPEMatching(regex, fullSearch=False):		def getCPEMatching(regex, fullSearch=False):
	lst=list(colCPE.find({"id": {"$regex": regex}}))		lst=list(colCPE.find({"id": {"$regex": regex}}))
	if fullSearch: lst.extend(colCPEOTHER.find({"id": {"$regex": regex}}))		if fullSearch: lst.extend(colCPEOTHER.find({"id": {"$regex": regex}}))
	return lst		return lst

	def getAlternativeCPE(id):
	return sanitize(colCPEOTHER.find_one({"id": id}))

	def getFreeText(text):		def getFreeText(text):
	return [x["obj"] for x in db.command("text", "cves", search=text)["results"]]		return [x["obj"] for x in db.command("text", "cves", search=text)["results"]]

			def getCAPECFor(cwe):
			return sanitize(colCAPEC.find({'related_weakness': {'$in': [cwe]}}))

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

sbin/db_cpe_browser.py

+4 −6

Original line number	Original line	Diff line number	Diff line
	@@ -21,26 +21,24 @@ from redis import exceptions as redisExceptions

	from lib.Config import Configuration		from lib.Config import Configuration
	from lib.Toolkit import pad		from lib.Toolkit import pad
			import lib.DatabaseLayer as db

	argParser = argparse.ArgumentParser(description='CPE entries importer in Redis cache')		argParser = argparse.ArgumentParser(description='CPE entries importer in Redis cache')
	argParser.add_argument('-v', action='store_true', default=False, help='Verbose logging')		argParser.add_argument('-v', action='store_true', default=False, help='Verbose logging')
	argParser.add_argument('-o', action='store_true', default=False, help='Import cpeother database in Redis cache')		argParser.add_argument('-o', action='store_true', default=False, help='Import cpeother database in Redis cache')
	args = argParser.parse_args()		args = argParser.parse_args()

	# connect to db
	db = Configuration.getMongoConnection()

	if args.o:		if args.o:
	cpe = db.cpeother		cpe = db.getAlternativeCPEs()
	else:		else:
	cpe = db.cpe		cpe = db.getCPEs()

	try:		try:
	r = Configuration.getRedisVendorConnection()		r = Configuration.getRedisVendorConnection()
	except:		except:
	sys.exit(1)		sys.exit(1)

	for e in cpe.find({}):		for e in cpe:
	try:		try:
	if args.o is not True:		if args.o is not True:
	prefix = 'cpe_2_2'		prefix = 'cpe_2_2'

sbin/db_mgmt_cwe.py

+2 −1

Original line number	Original line	Diff line number	Diff line
	@@ -103,10 +103,11 @@ f = open(os.path.join(tmpdir, 'cwec_v2.8.xml'))
	parser.parse(f)		parser.parse(f)
	cweList=[]		cweList=[]
	for cwe in progressbar(ch.cwe):		for cwe in progressbar(ch.cwe):
			cwe['description_summary']=cwe['description_summary'].replace("\t\t\t\t\t", " ")
	if args.v:		if args.v:
	print (cwe)		print (cwe)
	cweList.append(cwe)		cweList.append(cwe)
	db.cweBulkUpdate(cweList)		db.bulkUpdate('cwe', cweList)

	#update database info after successful program-run		#update database info after successful program-run
	db.setColUpdate('cwe', lastmodified)		db.setColUpdate('cwe', lastmodified)

sbin/db_mgmt_exploitdb.py

+6 −13

Original line number	Original line	Diff line number	Diff line
	@@ -16,7 +16,7 @@ runPath = os.path.dirname(os.path.realpath(__file__))
	sys.path.append(os.path.join(runPath, ".."))		sys.path.append(os.path.join(runPath, ".."))

	from lib.Config import Configuration		from lib.Config import Configuration
	import lib.DatabaseLayer as dbLayer		import lib.DatabaseLayer as db

	import csv		import csv
	import argparse		import argparse
	@@ -25,11 +25,6 @@ import argparse
	exploitdburl = Configuration.getexploitdbDict()		exploitdburl = Configuration.getexploitdbDict()
	tmppath = Configuration.getTmpdir()		tmppath = Configuration.getTmpdir()

	# connect to db
	db = Configuration.getMongoConnection()
	exploitdb = db.exploitdb
	info = db.info

	argparser = argparse.ArgumentParser(description='Populate/update the exploitdb ref database')		argparser = argparse.ArgumentParser(description='Populate/update the exploitdb ref database')
	argparser.add_argument('-v', action='store_true', help='verbose output', default=False)		argparser.add_argument('-v', action='store_true', help='verbose output', default=False)
	args = argparser.parse_args()		args = argparser.parse_args()
	@@ -40,7 +35,7 @@ try:
	except:		except:
	sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(exploitdburl))		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(exploitdburl))

	i = dbLayer.getLastModified('exploitdb')		i = db.getLastModified('exploitdb')
	if i is not None:		if i is not None:
	if f.headers['last-modified'] == i:		if f.headers['last-modified'] == i:
	print("Not modified")		print("Not modified")
	@@ -54,16 +49,14 @@ with open(csvfile, 'wb') as fp:
	shutil.copyfileobj(f, fp)		shutil.copyfileobj(f, fp)
	fp.close()		fp.close()

	bulk = exploitdb.initialize_ordered_bulk_op()		exploits=[]

	with open(csvfile, newline='') as csvtoparse:		with open(csvfile, newline='') as csvtoparse:
	exploitcsv = csv.DictReader(csvtoparse, delimiter=',')		exploitcsv = csv.DictReader(csvtoparse, delimiter=',')
	for row in exploitcsv:		for row in exploitcsv:
			exploits.append(row)
	if args.v:		if args.v:
	print ("{} ({}) Imported".format(row['id'],row['description']))		print ("{} ({}) Imported".format(row['id'],row['description']))
	bulk.find({'id': row['id']}).upsert().update({"$set": {'description': row['description'], 'type': row['type'], 'date': row['date'], 'port': row['port'], 'author': row['author'], 'file': row['file'], 'platform': row['platform'], 'id': row['id']}})		db.bulkUpdate('exploitdb', exploits)

	bulk.execute()

	# Update last-modified		# Update last-modified
	dbLayer.setColUpdate('exploitdb', f.headers['last-modified'])		db.setColUpdate('exploitdb', f.headers['last-modified'])

Admin message