More database layer abstraction (18382311) · Commits · Zhang Yushan / cve-search

lib/DatabaseLayer.py

+33 −0

Original line number	Diff line number	Diff line
		@@ -54,6 +54,18 @@ def updateCVE(cve):
		def updateMSBulletin(ms):
		colMSBULLETIN.update({"id": ms['id']},{"$set": ms}, upsert=True)

		def cpeBulkUpdate(cpelist):
		bulk=colCPE.initiate_ordered_bulk_op()
		for x in cpelist:
		# To check: Can we not just update({'$set': x})
		if x['references']:
		bulk.find({'id': x['id']}).upsert().update({'$set': {'title': x['title'], 'cpe_2_2': x['name'], 'references': x['references']}})
		else:
		bulk.find({'id': x['id']}).upsert().update({'$set': {'title': x['title'], 'cpe_2_2': x['name']}})

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

		def dropCollection(col):
		return db[col].drop()

		@@ -66,6 +78,24 @@ def cvesForCPE(cpe):
		return sanitize(colCVE.find({"vulnerable_configuration": {"$regex": cpe}}).sort("Modified", -1))

		# User Functions
		def addUser(user, pwd, admin=False):
		if admin:
		colUSERS.insert({'username':user, 'password':pwd, 'master': True})
		else:
		colUSERS.insert({'username':user, 'password':pwd})

		def changePassword(user, pwd):
		colUSERS.update({'username': user}, {'$set': {'password': pwd}})

		def deleteUser(user):
		colUSERS.remove({'username': user})

		def setAdmin(user, admin):
		if admin:
		colUSERS.update({'username': user}, {'$set': {'master': True}})
		else:
		colUSERS.update({'username': user}, {'$unset': {'master': ""}})

		def seenCVEs(user):
		data = colSEEN.find_one({"user": user})
		if not data:
		@@ -108,6 +138,9 @@ def getCVEs(limit=False, query=[], skip=0, cves=None):
		cve=colCVE.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)
		return sanitize(cve)

		def getCVEsNewerThan(dt):
		return sanitize(getCVEs(query={'last-modified': {'$gt': dt}}))

		def getCVEIDs(limit=-1):
		return [x["id"] for x in colCVE.find().limit(limit).sort("Modified", -1)]

sbin/db_mgmt_admin.py

+9 −11

Original line number	Diff line number	Diff line
		@@ -35,10 +35,8 @@ argParser.add_argument('-p', help='Promote account to master', default=False)
		argParser.add_argument('-d', help='Demote account to normal user', default=False)
		args = argParser.parse_args()

		# connect to db
		db = Configuration.getMongoConnection()
		collection = db.mgmt_users

		# vars
		col = "mgmt_users"
		rounds = 8000
		saltLength = 10
		exits = {'userInDb': 'User already exists in database',
		@@ -92,19 +90,19 @@ try:
		if dbLayer.userExists(username):
		sys.exit(exits['userInDb'])
		# set master if db is empty
		if(collection.count() > 0):
		if dbLayer.getSize(col) > 0:
		masterLogin()
		password = promptNewPass()
		collection.insert({'username': username, 'password': password})
		dbLayer.addUser(username, password)
		else:
		password = promptNewPass()
		collection.insert({'username': username, 'password': password, 'master': True})
		dbLayer.addUser(username, password, admin=True)
		sys.exit("User added")
		elif args.c:
		username = args.c
		verifyPass(getpass.getpass("Old password:"), username)
		password = promptNewPass()
		collection.update({'username': username}, {'$set': {'password': password}})
		dbLayer.changePassword(username, password)
		sys.exit("Password updated")
		elif args.r:
		username = args.r
		@@ -112,7 +110,7 @@ try:
		sys.exit(exits['userNotInDb'])
		masterLogin()
		isLastAdmin(username)
		collection.remove({'username': username})
		dbLayer.deleteUser(username)
		sys.exit('User removed from database')
		elif args.p:
		username = args.p
		@@ -120,7 +118,7 @@ try:
		sys.exit(exits['userNotInDb'])
		masterLogin()
		# promote
		collection.update({'username': username}, {'$set': {'master': True}})
		dbLayer.setAdmin(username, True)
		sys.exit('User promoted')
		elif args.d:
		username = args.d
		@@ -129,7 +127,7 @@ try:
		masterLogin()
		isLastAdmin(username)
		# demote
		collection.update({'username': username}, {'$unset': {'master': ""}})
		dbLayer.setAdmin(username, False)
		sys.exit('User demoted')

		except pymongo.errors.ConnectionFailure:

sbin/db_mgmt_cpe_dictionary.py

+9 −17

Original line number	Diff line number	Diff line
		@@ -28,7 +28,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 db

		class CPEHandler(ContentHandler):
		def __init__(self):
		@@ -70,10 +70,6 @@ class CPEHandler(ContentHandler):

		# dict
		cpedict = Configuration.getCPEDict()
		# connect to db
		db = Configuration.getMongoConnection()
		cpe = db.cpe
		info = db.info

		# make parser
		parser = make_parser()
		@@ -84,23 +80,19 @@ try:
		f = Configuration.getFile(cpedict)
		except:
		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(cpedict))
		i = info.find_one({'db': 'cpe'})
		i = db.getLastModified('cpe')
		if i is not None:
		if f.headers['last-modified'] == i['last-modified']:
		if f.headers['last-modified'] == i:
		print("Not modified")
		sys.exit(0)
		# parse xml and store in database
		parser.parse(f)
		bulk = cpe.initialize_ordered_bulk_op()
		cpeList=[]
		for x in progressbar(ch.cpe):
		name = toStringFormattedCPE(x['name'])
		oldCPE = x['name']
		title = x['title'][0]
		if x['references']:
		bulk.find({'id': name}).upsert().update({"$set":{'title': title, 'cpe_2_2':oldCPE, 'references': x['references']}})
		else:
		bulk.find({'id': name}).upsert().update({"$set":{'title': title, 'cpe_2_2':oldCPE}})
		bulk.execute()
		x['id']= toStringFormattedCPE(x['name'])
		x['title']=x['title'][0]
		cpeList.append(x)
		db.cpeBulkUpdate(cpeList)

		#update database info after successful program-run
		info.update({'db': 'cpe'}, {"$set": {'last-modified': f.headers['last-modified']}}, upsert=True)
		db.setColUpdate('cpe', f.headers['last-modified'])

sbin/db_mgmt_cpe_other_dictionary.py

+10 −14

Original line number	Diff line number	Diff line
		@@ -35,16 +35,12 @@ import urllib

		from lib.ProgressBar import progressbar
		from lib.Config import Configuration
		import lib.DatabaseLayer as db

		# connect to db
		db = Configuration.getMongoConnection()
		cpe = db.cpe
		cpeother = db.cpeother
		cve = db.cves
		info = db.info
		# get dates
		icve = db.getLastModified('cve')
		icpeo = db.getLastModified('cpeother')

		icve = info.find_one({'db': 'cve'})
		icpeo = info.find_one({'db': 'cpeother'})
		# check modification date
		date = False
		if icve is not None and icpeo is not None:
		@@ -58,9 +54,9 @@ if icve is not None and icpeo is not None:
		# only get collection of new CVE's
		collections = []
		if date:
		collections = cve.find({'last-modified': {'$gt': icve['last-modified']}})
		db.getCVEsNewerThan(icve['last-modified'])
		else:
		collections = cve.find({})
		db.getCVEs()
		# check cpes for cves and parse and store missing cpes in cpeother
		batch = []

		@@ -72,9 +68,9 @@ if not col:

		for item in progressbar(col):
		for cpeentry in item['vulnerable_configuration']:
		checkdup = cpeother.find(({'id': cpeentry}))
		checkdup = db.getAlternativeCPE(cpeentry)
		if checkdup.count() <= 0:
		entry = cpe.find(({'id': cpeentry}))
		entry = db.getCPE(cpeentry)
		if entry.count() <= 0:
		title = cpeentry
		title = title[10:]
		@@ -86,7 +82,7 @@ for item in progressbar(col):
		title = title.title()
		batch.append({'id': cpeentry, 'title': title})
		if len(batch) != 0:
		cpeother.insert(batch)
		db.cpeotherBulkInsert(batch)

		#update database info after successful program-run
		info.update({'db': 'cpeother'}, {"$set": {'last-modified': icve['last-modified']}}, upsert=True)
		db.setColUpdate('cpeother', icve['last-modified'])

Admin message