more db layer abstraction + bugfix import/export (0a2d147d) · Commits · Zhang Yushan / cve-search

lib/DatabaseLayer.py

+39 −0

Original line number	Diff line number	Diff line
		@@ -40,6 +40,9 @@ def sanitize(x):
		def ensureIndex(col, field):
		db[col].ensure_index(field)

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

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

		@@ -229,9 +232,45 @@ def getDBStats():
		def getWhitelist():
		return sanitize(colWHITELIST.find())

		def isInWhitelist(cpe):
		return True if colWHITELIST.find({'id': cpe}).count()>0 else False

		def addToWhitelist(cpe, type, comments=None):
		if comments:
		colWHITELIST.insert({'id': cpe, 'type': type, 'comments': comments})
		else:
		colWHITELIST.insert({'id': cpe, 'type': type})

		def removeFromWhitelist(cpe):
		colWHITELIST.remove({'id': cpe})

		def updateWhitelist(oldCPE, newCPE, type, comments=None):
		if comments:
		colWHITELIST.update({'id': oldCPE}, {'id': newCPE, 'type': type, 'comments': comments})
		else:
		colWHITELIST.update({'id': oldCPE}, {'id': newCPE, 'type': type})

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

		def isInBlacklist(cpe):
		return True if colBLACKLIST.find({'id': cpe}).count()>0 else False

		def addToBlacklist(cpe, type, comments=None):
		if comments:
		colBLACKLIST.insert({'id': cpe, 'type': type, 'comments': comments})
		else:
		colBLACKLIST.insert({'id': cpe, 'type': type})

		def removeFromBlacklist(cpe):
		colBLACKLIST.remove({'id': cpe})

		def updateBlacklist(oldCPE, newCPE, type, comments=None):
		if comments:
		colBLACKLIST.update({'id': oldCPE}, {'id': newCPE, 'type': type, 'comments': comments})
		else:
		colBLACKLIST.update({'id': oldCPE}, {'id': newCPE, 'type': type})

		def getRules(list):
		if list.lower()=='whitelist':
		col=colWHITELIST

lib/cpelist.py

+37 −50

Original line number	Diff line number	Diff line
		@@ -18,32 +18,26 @@ sys.path.append(os.path.join(runPath, ".."))
		import re

		from lib.Toolkit import toStringFormattedCPE
		import lib.DatabaseLayer as db

		class CPEList:

		def __init__(self, collection, args):
		self.collection = collection
		self.collection = collection.title()
		self.args = args

		# check if there are items in the collection
		def countItems(self):
		return self.collection.count()
		return db.getSize("mgmt_"+self.collection.lower())

		# check if a cpe is in the list
		def check(self, cpe):
		try:
		cpeListElement = {'id': cpe}
		amount = self.collection.find(cpeListElement).count()
		return amount
		except Exception:
		print("Error connecting to the database")
		sys.exit()
		return getattr(db,"isIn"+self.collection)(cpe)

		# insert to database
		def insert(self, cpe, cpeType):
		try:
		# split comments from cpe
		if '#' in cpe:
		comments = cpe.split('#')
		del comments[0]
		cpeID = cpe.split('#')[0]
		@@ -52,12 +46,8 @@ class CPEList:
		# check format
		if cpeID:
		# already in db?
		if self.check(cpeID) == 0:
		if '#' in cpe:
		cpeListElement = {'id': cpeID, 'type':cpeType, 'comments': comments}
		else:
		cpeListElement = {'id': cpeID, 'type':cpeType}
		self.collection.insert(cpeListElement)
		if not self.check(cpeID):
		getattr(db, "addTo"+self.collection)(cpeID, cpeType, comments)
		return True
		return False
		except Exception as ex:
		@@ -71,13 +61,13 @@ class CPEList:
		# translate cpe
		if toStringFormattedCPE(cpe): cpe = toStringFormattedCPE(cpe)
		# check if the cpe is in the list
		amount = self.check(cpe)
		if amount > 0:
		cpeListElement = {'id': cpe}
		self.collection.remove(cpeListElement)
		return amount
		if self.check(cpe):
		getattr(db, "removeFrom"+self.collection)(cpe)
		return True
		else:
		return False
		except Exception as ex:
		print("Error removing item from database: {:d}".format(ex))
		print("Error removing item from database: %s"%(ex))
		sys.exit()

		def update(self, cpeOld, cpeNew, cpeType):
		@@ -89,35 +79,29 @@ class CPEList:
		cpeNew = toStringFormattedCPE(cpeNew)
		if cpeOld and cpeNew:
		# already in db?
		if self.check(cpeOld.split('#')[0]) != 0:
		if '#' in cpeNew:
		# there are extra comments
		if self.check(cpeOld.split('#')[0]):
		cpeID = cpeNew.split('#')[0]
		cpeID.strip()
		# allow multiple comments
		# comments
		comments = cpeNew.split('#')
		del comments[0]
		cpeListElement = {'id': cpeID, 'comments': comments, 'type':cpeType}
		else:
		cpeListElement = {'id': cpeNew, 'type':cpeType}
		cpeDeleteElement = {'id': cpeOld.split('#')[0]}
		self.collection.update(cpeDeleteElement, cpeListElement, upsert=False, multi=False)
		getattr(db, "update"+self.collection)(cpeOld.split('#')[0], cpeID, cpeType, comments)
		return True
		return False
		except Exception as ex:
		print(ex)
		print("Error updating item in database: {:d}".format(ex))
		print("Error updating item in database: %s"%(ex))
		sys.exit()

		# drop the collection
		def dropCollection(self):
		try:
		count = self.countItems()
		self.collection.drop()
		db.drop("mgmt_"+self.collection.lower())
		if self.args.v:
		print("collection of {:d} items dropped".format(count))
		print("collection of %s items dropped"%(count))
		except Exception as ex:
		print("Error dropping the database: {:d}".format(ex))
		print("Error dropping the database: %s"%(ex))
		sys.exit()

		# import a file that represents the cpe list
		@@ -126,11 +110,14 @@ class CPEList:
		# read each line from the import file and regex them to a cpe format
		try:
		for line in importFile:
		print(self.args.t)
		if self.insert(line,self.args.t):
		try:
		cpe, cpeType=line.strip().split("\|")
		if self.insert(cpe, cpeType):
		count += 1
		except:
		continue
		if self.args.v:
		print("{:d} products added to the list".format(count))
		print("%s products added to the list"%(count))
		except IOError:
		print('Could not open the file')
		sys.exit()
		@@ -140,7 +127,7 @@ class CPEList:
		count = 0
		# check if file exists already
		if not os.path.exists(exportFile) or self.args.f:
		listed = self.collection.find()
		listed = getattr(db, "get"+self.collection)()
		export = open(exportFile, 'w')
		for listedID in listed:
		count += 1
		@@ -151,10 +138,10 @@ class CPEList:
		# separate the comments
		for comment in comments:
		commentString = commentString + '#' + comment
		export.write(listedID['id'] + commentString + '\n')
		export.write(listedID['id'] + commentString + "\|" + listedID["type"] + '\n')
		export.close()
		if self.args.v:
		print("{:d} listed items exported".format(count))
		print("%s listed items exported"%(count))
		else:
		print("file already exists")

		@@ -191,7 +178,7 @@ class CPEList:
		if self.insert(cpeID,self.args.t):
		count += 1
		if self.args.v:
		print("{:d} products added to the list".format(count))
		print("%s products added to the list"%(count))
		elif self.args.r or self.args.R:
		# get list of cpe's to remove
		if self.args.r:
		@@ -204,4 +191,4 @@ class CPEList:
		amount = self.remove(cpeID)
		count += amount
		if self.args.v:
		print("{:d} products removed from the list".format(count))
		print("%s products removed from the list"%(count))

sbin/db_blacklist.py

+3 −4

Original line number	Diff line number	Diff line
		@@ -33,11 +33,10 @@ argparser.add_argument('-f', action='store_true', help='Force')
		argparser.add_argument('-v', action='store_true', help='Verbose')
		args = argparser.parse_args()

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

		# Variables
		collection = "blacklist"

		# Functions
		def importBlacklist(importFile):
		oList = CPEList(collection, args)
		oList.importList(importFile)

sbin/db_whitelist.py

+2 −4

Original line number	Diff line number	Diff line
		@@ -33,10 +33,8 @@ argparser.add_argument('-f', action='store_true', help='Force')
		argparser.add_argument('-v', action='store_true', help='Verbose')
		args = argparser.parse_args()

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

		# Variables
		collection = "whitelist"

		def importWhitelist(importFile):
		oList = CPEList(collection, args)

Admin message