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

lib/CVEs.py

+3 −8

Original line number	Diff line number	Diff line
		@@ -21,6 +21,7 @@ import math

		from lib.Config import Configuration
		from lib.Toolkit import exploitabilityScore,impactScore
		import lib.DatabaseLayer as db

		class last:

		@@ -40,9 +41,6 @@ class last:

		if rankinglookup:
		self.ranking = connectdb['ranking']
		if namelookup:
		self.cpeOther = connectdb['cpeother']
		self.cpe = connectdb['cpe']
		if vfeedlookup:
		self.vfeed = connectdb['vfeed']
		if capeclookup:
		@@ -67,14 +65,11 @@ class last:
		return ref

		def getcpe(self, cpeid=None):

		if not(self.namelookup):
		return cpeid

		e = self.cpe.find_one({'id': cpeid})

		e = db.getCPE(cpeid)
		if e is None:
		e = self.cpeOther.find_one({'id': cpeid})
		e = db.getAlternativeCPE(cpeid)
		if e is None:
		return cpeid
		if 'id' in e:

lib/DatabaseLayer.py

+47 −10

Original line number	Diff line number	Diff line
		@@ -11,35 +11,72 @@

		# imports
		from lib.Config import Configuration as conf
		import pymongo

		# Variables
		db=conf.getMongoConnection()
		colCVE= db['cves']
		colCPE= db['cpe']
		colCPEOTHER= db['cpeother']
		colWHITELIST=db['mgmt_whitelist']
		colBLACKLIST=db['mgmt_blacklist']
		colUSERS= db['mgmt_users']
		colSEEN= db['mgmt_seen']

		# Functions
		def sanitize(x):
		if type(x)==pymongo.cursor.Cursor:
		x=list(x)
		if type(x)==list:
		for y in x: sanitize(y)
		if "_id" in x: x.pop("_id")
		return x

		# DB Functions

		# API Functions
		def cvesForCPE(cpe):
		col=db['cves']
		if not cpe: return []
		cves=list(col.find({"vulnerable_configuration": {"$regex": cpe}}).sort("Modified", -1))
		for cve in cves:
		cve.pop("_id")
		return cves
		return sanitize(colCVE.find({"vulnerable_configuration": {"$regex": cpe}}).sort("Modified", -1))

		# User Functions
		def seenCVEs(user):
		col=db['mgmt_seen']
		data = col.find_one({"user": user})
		data = colSEEN.find_one({"user": user})
		if not data:
		col.insert({"user": user, "seen_cves": []})
		colSEEN.insert({"user": user, "seen_cves": []})
		return []
		else:
		return data['seen_cves']

		def addSeenCVEs(user, CVEs):
		col=db['mgmt_seen']
		if type(CVEs) == str: CVEs=[CVEs]
		if type(CVEs) == list:
		seen=list(set(CVEs)-set(seenCVEs(user)))
		if seen:
		col.update({"user": user},{"$addToSet": {"seen_cves": { "$each": seen}}})
		colSEEN.update({"user": user},{"$addToSet": {"seen_cves": { "$each": seen}}})

		# Query Functions
		def getCVEs(limit=-1, query=[], skip=0):
		if type(query) == dict: query=[query]
		if len(query) == 0:
		cves=colCVE.find().sort("Modified", -1).limit(limit).skip(skip)
		elif len(query) == 1:
		cves=colCVE.find(query[0]).sort("Modified", -1).limit(limit).skip(skip)
		else:
		cves=colCVE.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)
		return sanitize(cves)

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

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

		def getUsers():
		return sanitize(colUSERS.find())

		def getWhitelist():
		return sanitize(colWHITELIST.find())

		def getBlacklist():
		return sanitize(colBLACKLIST.find())

lib/Query.py

+3 −4

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

		from lib.Config import Configuration

		import lib.DatabaseLayer as dbLayer

		db = Configuration.getMongoConnection()
		collection = db.cves
		@@ -48,7 +48,7 @@ def findranking(cpe=None, loosy=True):
		return result

		def lookupcpe(cpeid=None):
		e = db.cpe.find_one({'id': cpeid})
		e = dbLayer.getCPE(cpeid)
		if e is None:
		return cpeid
		if 'id' in e:
		@@ -57,8 +57,7 @@ def lookupcpe(cpeid=None):

		def lastentries(limit=5, namelookup=False, rankinglookup=True):
		entries = []
		for item in collection.find({}).sort("Modified", -1).limit(limit):
		item.pop('_id')
		for item in dbLayer.getCVEs(limit):
		if not namelookup and rankinglookup is not True:
		entries.append(item)
		else:

lib/User.py

+4 −7

Original line number	Diff line number	Diff line
		@@ -18,16 +18,13 @@ runPath = os.path.dirname(os.path.realpath(__file__))
		from flask.ext.login import UserMixin

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

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


		# Exception
		class UserNotFoundError(Exception):
		pass


		# Class
		class User(UserMixin):

		'''Simple User class'''
		@@ -36,7 +33,7 @@ class User(UserMixin):
		USERS = {"_dummy_": "_dummy_"}
		else:
		USERS = {}
		for user in collection.find({}):
		for user in db.getUsers():
		USERS[user['username']] = user['password']

		def __init__(self, id):

web/index.py

+6 −25

Original line number	Diff line number	Diff line
		@@ -83,13 +83,6 @@ def getBrowseList(vendor):
		return result


		def getWhitelist():
		collection = db.mgmt_whitelist
		whitelist = list(collection.find())
		for x in whitelist:x.pop("_id")
		return whitelist


		def getWhitelistRules():
		collection = db.mgmt_whitelist
		whitelist = collection.find({'type':'cpe'}).distinct('id')
		@@ -139,13 +132,6 @@ def seen_mark(cve):
		for c in cve:
		if c["id"] in seen: cve[cve.index(c)]['seen'] = 'yes'

		def getBlacklist():
		collection = db.mgmt_blacklist
		blacklist = list(collection.find())
		for x in blacklist:x.pop("_id")
		return blacklist


		def getBlacklistRules():
		collection = db.mgmt_blacklist
		blacklist = collection.find({'type':'cpe'}).distinct('id')
		@@ -270,12 +256,7 @@ def filter_logic(blacklist, whitelist, unlisted, timeSelect, startDate, endDate,
		query.append({timeTypeSelect: {'$gt': startDate, '$lt': endDate}})
		if timeSelect == "outside":
		query.append({'$or': [{timeTypeSelect: {'$lt': startDate}}, {timeTypeSelect: {'$gt': endDate}}]})
		if len(query) == 0:
		cve = collection.find().sort("Modified", -1).limit(limit).skip(skip)
		elif len(query) == 1:
		cve = collection.find(query[0]).sort("Modified", -1).limit(limit).skip(skip)
		else:
		cve = collection.find({'$and': query}).sort("Modified", -1).limit(limit).skip(skip)
		cve=dbLayer.getCVEs(limit=limit, skip=skip, query=query)
		# marking relevant records
		if whitelist == "on":
		cve = whitelist_mark(cve)
		@@ -570,7 +551,7 @@ def whitelistDrop():
		@app.route('/admin/whitelist')
		@login_required
		def whitelistView():
		return render_template('list.html', rules=getWhitelist(), status=["default", "none"], listType="Whitelist")
		return render_template('list.html', rules=dbLayer.getWhitelist(), status=["default", "none"], listType="Whitelist")


		@app.route('/admin/addToList')
		@@ -580,7 +561,7 @@ def listAdd():
		cpeType = request.args.get('type')
		lst = request.args.get('list')
		status = ["added", "success"] if addCPEToList(cpe, lst, cpeType) else ["already_exists", "info"]
		returnList = getWhitelist() if lst=="whitelist" else getBlacklist()
		returnList = dbLayer.getWhitelist() if lst=="whitelist" else dbLayer.getBlacklist()
		return jsonify({"status":status, "rules":returnList, "listType":lst.title()})


		@@ -597,7 +578,7 @@ def listRemove():
		status = ["removed", "success"] if (result > 0) else ["already_removed", "info"]
		else:
		status = ["invalid_url", "error"]
		returnList = getWhitelist() if lst=="whitelist" else getBlacklist()
		returnList = dbLayer.getWhitelist() if lst=="whitelist" else dbLayer.getBlacklist()
		return jsonify({"status":status, "rules":returnList, "listType":lst.title()})


		@@ -616,7 +597,7 @@ def listEdit():
		status = ["update_failed", "error"]
		else:
		status = ["invalid_url", "error"]
		returnList = list(getWhitelist()) if lst=="whitelist" else list(getBlacklist())
		returnList = list(dbLayer.getWhitelist()) if lst=="whitelist" else list(dbLayer.getBlacklist())
		return jsonify({"rules":returnList, "status":status, "listType":lst})


		@@ -660,7 +641,7 @@ def blacklistDrop():
		@app.route('/admin/blacklist')
		@login_required
		def blacklistView():
		return render_template('list.html', rules=getBlacklist(), status=["default", "none"], listType="Blacklist")
		return render_template('list.html', rules=dbLayer.getBlacklist(), status=["default", "none"], listType="Blacklist")


		@app.route('/admin/listmanagement/add')

Admin message