VIA4 intake (def16a65) · Commits · Zhang Yushan / cve-search

lib/Database.py

+8 −0

Original line number	Diff line number	Diff line
		@@ -109,6 +109,10 @@ class Database(metaclass=Singleton):
		'data': data}
		return data

		def db_ensureIndex(self, collection, field):
		self.db[collection].ensure_index(field)


		#######################
		# Blacklist/Whitelist #
		#######################
		@@ -373,6 +377,10 @@ class Database(metaclass=Singleton):
		def via4_info(self):
		return self._getInfo(self.colVIA4)

		def via4_setMetadata(self, field, data):
		self.colINFO.update({"db": self.colVIA4.name},
		{"$set": {field: data}}, upsert=True)

		def via4_setUpdate(self, date):
		self._setUpdate(self.colVIA4, date)

+15 −0

Original line number	Diff line number	Diff line
		@@ -40,6 +40,9 @@ class DatabaseLayer(metaclass=Singleton):
		def db_info(self, include_admin=False):
		return self.db.db_getStats(include_admin)

		def ensure_index(self, collection, field):
		self.db.db_ensureIndex(collection, field)

		####################
		# Black-/Whitelist #
		####################
		@@ -392,6 +395,18 @@ class VIA4s:
		via4 = self.db.via4_get(cveid)
		return via4 if via4 else None

		def searchables(self, data=None):
		if data:
		self.db.via4_setMetadata('searchables', data)
		else:
		return self.db.via4_info().get('searchables', [])

		def sources(self, data=None):
		if data:
		self.db.via4_setMetadata('sources', data)
		else:
		return self.db.via4_info().get('sources', [])

		def updated(self, date=None):
		if date:
		self.db.via4_setUpdate(date)

+13 −6

Original line number	Diff line number	Diff line
		@@ -236,14 +236,21 @@ class VIA4:
		return values
		return VIA4(**data) if isinstance(data, dict) else data

		def dict(self):
		data = {}
		for key, val in self.__dict__.items():
		def _dig(self, data):
		if isinstance(data, dict):
		for key, val in data.items():
		if key.startswith("__"): continue
		if isinstance(val, VIA4): val = val.dict()
		data[key] = val
		data[key] = self._dig(val)
		if isinstance(data, VIA4):
		data = self._dig(data.__dict__)
		if isinstance(data, list):
		for i, item in enumerate(data):
		data[i] = self._dig(item)
		return data

		def dict(self):
		return self._dig(self.__dict__)

		@classmethod
		def fromDict(cls, data):
		return cls(**data)

+6 −6

Original line number	Diff line number	Diff line
		@@ -16,11 +16,13 @@ sys.path.append(os.path.join(runPath, ".."))

		from pymongo import TEXT

		import lib.DatabaseLayer as dbLayer
		from lib.DatabaseLayer2 import DatabaseLayer

		db = DatabaseLayer()

		def setIndex(col, field, printSuccess = True):
		try:
		dbLayer.ensureIndex(col, field)
		db.ensure_index(col, field)
		if printSuccess:
		print('[+]Success to create index %s on %s' % (field, col))
		except Exception as e:
		@@ -38,7 +40,5 @@ setIndex('mgmt_whitelist', 'id')
		setIndex('mgmt_blacklist', 'id')
		setIndex('capec', 'related_weakness')

		via4 = dbLayer.getInfo('via4')
		if via4:
		for index in via4.get('searchables', []):
		for index in db.VIA4.searchables():
		setIndex('via4', index, False)

+10 −8

Original line number	Diff line number	Diff line
		@@ -18,7 +18,8 @@ sys.path.append(os.path.join(runPath, ".."))
		from dateutil.parser import parse as parse_datetime

		from lib.Config import Configuration
		import lib.DatabaseLayer as db
		from lib.DatabaseLayer2 import DatabaseLayer
		from lib.Objects import VIA4

		# To Do: Implement REDIS

		@@ -37,9 +38,10 @@ try:
		except:
		sys.exit("Cannot open url %s. Bad URL or not connected to the internet?"%(Configuration.getFeedURL("via4")))

		db = DatabaseLayer()
		# check modification date
		lastmodified = parse_datetime(r.headers['last-modified'], ignoretz=True)
		i=db.getLastModified("via4")
		i = db.VIA4.updated()
		if i is not None:
		if lastmodified == i:
		print("Not modified")
		@@ -47,10 +49,10 @@ if i is not None:

		data = json.loads(f.read().decode('utf-8'))
		cves = data['cves']
		bulk = [dict(val, id=key) for key, val in cves.items() if key]
		db.bulkUpdate('via4', bulk)
		db.setColInfo('via4', 'sources', data['metadata']['sources'])
		db.setColInfo('via4', 'searchables', data['metadata']['searchables'])
		bulk = [VIA4.fromDict(dict(val, id=key)) for key, val in cves.items() if key]
		db.VIA4.upsert(bulk)
		db.VIA4.sources(data['metadata']['sources'])
		db.VIA4.searchables(data['metadata']['searchables'])

		#update database info after successful program-run
		db.setColUpdate('via4', lastmodified)
		db.VIA4.updated(lastmodified)