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

lib/DatabaseLayer.py

+20 −10

Original line number	Diff line number	Diff line
		@@ -26,6 +26,7 @@ colSEEN= db['mgmt_seen']
		colINFO= db['info']
		colVFEED= db['vfeed']
		colRANKING= db['ranking']
		colMSBULLETIN=db['ms']
		# Functions
		def sanitize(x):
		if type(x)==pymongo.cursor.Cursor:
		@@ -50,6 +51,15 @@ def updateCVE(cve):
		"cwe": cve['cwe'], "vulnerable_configuration": cve['vulnerable_configuration'],
		"vulnerable_configuration_cpe_2_2": cve['vulnerable_configuration_cpe_2_2'], 'last-modified': cve['Modified']}})

		def updateMSBulletin(ms):
		colMSBULLETIN.update({"id": ms['id']},{"$set": ms}, upsert=True)

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

		def getTableNames():
		return db.collection_names()

		# API Functions
		def cvesForCPE(cpe):
		if not cpe: return []

+4 −3

Original line number	Diff line number	Diff line
		@@ -19,6 +19,7 @@ import argparse

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

		argparser = argparse.ArgumentParser(description='populate/update d2sec exploit database')
		argparser.add_argument('-v', action='store_true', help='verbose output')
		@@ -104,9 +105,9 @@ try:
		f = Configuration.getFile(d2securl)
		except:
		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(d2securl))
		i = info.find_one({'db': 'd2sec'})
		i = dbLayer.getLastModified("d2sec")
		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
		@@ -120,4 +121,4 @@ for exploit in progressbar(ch.d2sec):
		bulk.execute()

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

+4 −8

Original line number	Diff line number	Diff line
		@@ -9,6 +9,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

		bulletinurl = "https://technet.microsoft.com/library/security/"

		@@ -23,18 +24,13 @@ def minimalist_xldate_as_datetime(xldate, datemode):
		msbulletinurl = Configuration.getMSBULLETINDict()
		tmppath = Configuration.getTmpdir()

		# connect to db
		db = Configuration.getMongoConnection()
		msbulletin = db.ms
		info = db.info

		try:
		f = Configuration.getFile(msbulletinurl)
		except:
		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(msbulletinurl))

		# check modification date
		i = info.find_one({'db': 'ms'})
		i=dbLayer.getInfo("ms")
		if i is not None:
		if f.headers['last-modified'] == i['last-modified']:
		print("Not modified")
		@@ -69,7 +65,7 @@ for rownum in range(sh.nrows-1):
		bulletin['software'].append(software)
		else:
		if bulletin :
		msbulletin.update({'id': bulletin['id']}, {'$set': bulletin}, upsert=True)
		dbLayer.updateMSBulletin(bulletin)
		bulletin = {}
		software = {}
		bulletin['Published'] = minimalist_xldate_as_datetime(row[0], 0).isoformat()
		@@ -93,7 +89,7 @@ for rownum in range(sh.nrows-1):
		prev_id = row[1]

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

+4 −7

Original line number	Diff line number	Diff line
		@@ -19,7 +19,7 @@ import argparse
		import time

		from lib.Config import Configuration

		import lib.DatabaseLayer as dbLayer

		runPath = os.path.dirname(os.path.realpath(__file__))

		@@ -65,14 +65,12 @@ if args.f and args.l:
		def nbelement(collection=None):
		if collection is None:
		collection = "cves"
		c = Configuration.getMongoConnection()
		return c[collection].count()
		return dbLayer.getSize(collection)

		def dropcollection(collection=None):
		if collection is None:
		return False
		c = Configuration.getMongoConnection()
		return c[collection].drop()
		return dbLayer.dropCollection(collection)

		def logging(message=None):
		if args.v:
		@@ -98,8 +96,7 @@ while (loop):
		dropcollection(collection=source['name'])
		logging( source['name'] + " dropped")
		if source['name'] is "cpeother":
		db = Configuration.getMongoConnection()
		if "cpeother" not in db.collection_names():
		if "cpeother" not in dbLayer.getTableNames():
		continue
		if source['name'] is not "redis-cache-cpe":
		message = 'Starting ' + source['name']