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

bin/search.py

+10 −11

Original line number	Diff line number	Diff line
		@@ -23,10 +23,7 @@ from bson import json_util

		from lib import CVEs
		from lib.Config import Configuration

		# connect to DB
		db = Configuration.getMongoConnection()
		collection = db.cves
		import lib.DatabaseLayer as db

		# init control variables
		csvOutput = 0
		@@ -53,7 +50,7 @@ argParser.add_argument('-a', default=False, action='store_true', help='Lookup CA
		argParser.add_argument('-v', type=str, help='vendor name to lookup in reference URLs')
		args = argParser.parse_args()
		vSearch = args.p
		cveSearch = args.c
		cveSearch = [x.upper() for x in args.c] if args.c else None
		vOutput = args.o
		vFreeSearch = args.f
		sLatest = args.l
		@@ -122,20 +119,22 @@ def printCVE(item):
		print(json.dumps(item, sort_keys=True, default=json_util.default))

		if cveSearch:
		for cveid in cveSearch:
		for item in collection.find({'id': cveid}).sort("Modified", sorttype):
		printCVE(item)
		for cveid in db.getCVEs(cves=cveSearch):
		printCVE(cveid)
		sys.exit(0)
		# Basic freetext search (in vulnerability summary).
		# Full-text indexing is more efficient to search across all CVEs.
		if vFreeSearch:
		for item in collection.find({'summary': {'$regex': re.compile(vFreeSearch, re.IGNORECASE)}}).sort("Modified", sorttype):
		try:
		for item in db.getFreeText(vFreeSearch):
		print(item)
		except:
		sys.exit("Free text search not enabled on the database!")
		sys.exit(0)

		# Search Product (best to use CPE notation, e.g. cisco:ios:12.2
		if vSearch:
		for item in collection.find({"vulnerable_configuration": {'$regex': vSearch}}).sort("Modified", sorttype):
		for item in db.cvesForCPE(vSearch):
		if csvOutput:
		# We assume that the vendor name is usually in the hostame of the
		# URL to avoid any match on the resource part

lib/DatabaseLayer.py

+12 −8

Original line number	Diff line number	Diff line
		@@ -25,6 +25,7 @@ colUSERS= db['mgmt_users']
		colSEEN= db['mgmt_seen']
		colINFO= db['info']
		colVFEED= db['vfeed']
		colRankking =db['ranking']
		# Functions
		def sanitize(x):
		if type(x)==pymongo.cursor.Cursor:
		@@ -86,15 +87,16 @@ def isSingleMaster(user):

		# Query Functions
		# Generic data
		def getCVEs(limit=-1, query=[], skip=0):
		def getCVEs(limit=False, query=[], skip=0, cves=None):
		if type(query) == dict: query=[query]
		if type(cves) == list: query.append({"id": {"$in": cves}})
		if len(query) == 0:
		cves=colCVE.find().sort("Modified", -1).limit(limit).skip(skip)
		cve=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)
		cve=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)
		cve=colCVE.find({"$and": query}).sort("Modified", -1).limit(limit).skip(skip)
		return sanitize(cve)

		def getCVEIDs(limit=-1):
		return [x["id"] for x in colCVE.find().limit(limit).sort("Modified", -1)]
		@@ -146,6 +148,9 @@ def getDBStats():
		def getWhitelist():
		return sanitize(colWHITELIST.find())

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

		def getRules(list):
		if list.lower()=='whitelist':
		col=colWHITELIST
		@@ -160,9 +165,8 @@ def getRules(list):
		rlist.extend(software)
		return rlist


		def getBlacklist():
		return sanitize(colBLACKLIST.find())
		def findRanking(cpe):
		return sanitize(colRANKING.find_one({'cpe': {'$regex': cpe}}))

		# Users
		def getUsers():

lib/Query.py

+5 −10

Original line number	Diff line number	Diff line
		@@ -17,29 +17,24 @@ 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
		import lib.DatabaseLayer as db

		rankinglookup = True


		def findranking(cpe=None, loosy=True):
		if cpe is None:
		return False
		r = db.ranking
		result = False
		if loosy:
		for x in cpe.split(':'):
		if x is not '':
		i = r.find_one({'cpe': {'$regex': x}})
		i = db.findRanking(cpe)
		if i is None:
		continue
		if 'rank' in i:
		result = i['rank']
		else:
		i = r.find_one({'cpe': {'$regex': cpe}})
		i = db.findRanking(cpe)
		print (cpe)
		if i is None:
		return result
		@@ -48,7 +43,7 @@ def findranking(cpe=None, loosy=True):
		return result

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

		def lastentries(limit=5, namelookup=False, rankinglookup=True):
		entries = []
		for item in dbLayer.getCVEs(limit):
		for item in db.getCVEs(limit):
		if not namelookup and rankinglookup is not True:
		entries.append(item)
		else:

Admin message