secator 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

secator/exporters/txt.py

@@ -1,24 +1,28 @@
 from secator.exporters._base import Exporter
+from secator.output_types import Info
 from secator.rich import console
 
 
 class TxtExporter(Exporter):
-    def send(self):
-        results = self.report.data['results']
-        txt_paths = []
+	def send(self):
+		results = self.report.data['results']
+		if not results:
+			return
+		txt_paths = []
 
-        for output_type, items in results.items():
-            items = [str(i) for i in items]
-            if not items:
-                continue
-            txt_path = f'{self.report.output_folder}/report_{output_type}.txt'
-            with open(txt_path, 'w') as f:
-                f.write('\n'.join(items))
-            txt_paths.append(txt_path)
+		for output_type, items in results.items():
+			items = [str(i) for i in items]
+			if not items:
+				continue
+			txt_path = f'{self.report.output_folder}/report_{output_type}.txt'
+			with open(txt_path, 'w') as f:
+				f.write('\n'.join(items))
+			txt_paths.append(txt_path)
 
-        if len(txt_paths) == 1:
-            txt_paths_str = txt_paths[0]
-        else:
-            txt_paths_str = '\n   • ' + '\n   • '.join(txt_paths)
+		if len(txt_paths) == 1:
+			txt_paths_str = txt_paths[0]
+		else:
+			txt_paths_str = '\n   • ' + '\n   • '.join(txt_paths)
 
-        console.print(f':file_cabinet: Saved TXT reports to {txt_paths_str}')
+		info = Info(f'Saved TXT reports to {txt_paths_str}')
+		console.print(info)

secator/hooks/gcs.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+from time import time
+
+from google.cloud import storage
+
+from secator.config import CONFIG
+from secator.runners import Task
+from secator.thread import Thread
+from secator.utils import debug
+
+
+GCS_BUCKET_NAME = CONFIG.addons.gcs.bucket_name
+ITEMS_TO_SEND = {
+	'url': ['screenshot_path']
+}
+
+
+def process_item(self, item):
+	if item._type not in ITEMS_TO_SEND.keys():
+		return item
+	if not GCS_BUCKET_NAME:
+		debug('skipped since addons.gcs.bucket_name is empty.', sub='hooks.gcs')
+		return item
+	to_send = ITEMS_TO_SEND[item._type]
+	for k, v in item.toDict().items():
+		if k in to_send and v:
+			path = Path(v)
+			if not path.exists():
+				continue
+			ext = path.suffix
+			blob_name = f'{item._uuid}_{k}{ext}'
+			t = Thread(target=upload_blob, args=(GCS_BUCKET_NAME, v, blob_name))
+			t.start()
+			self.threads.append(t)
+			setattr(item, k, f'gs://{GCS_BUCKET_NAME}/{blob_name}')
+	return item
+
+
+def upload_blob(bucket_name, source_file_name, destination_blob_name):
+	"""Uploads a file to the bucket."""
+	start_time = time()
+	storage_client = storage.Client()
+	bucket = storage_client.bucket(bucket_name)
+	blob = bucket.blob(destination_blob_name)
+	blob.upload_from_filename(source_file_name)
+	end_time = time()
+	elapsed = end_time - start_time
+	debug(f'in {elapsed:.4f}s', obj={'blob': 'CREATED', 'blob_name': destination_blob_name, 'bucket': bucket_name}, obj_after=False, sub='hooks.gcs', verbose=True)  # noqa: E501
+
+
+HOOKS = {
+	Task: {'on_item': [process_item]}
+}

secator/hooks/mongodb.py

@@ -6,7 +6,7 @@ from bson.objectid import ObjectId
 from celery import shared_task
 
 from secator.config import CONFIG
-from secator.output_types import OUTPUT_TYPES
+from secator.output_types import FINDING_TYPES
 from secator.runners import Scan, Task, Workflow
 from secator.utils import debug, escape_mongodb_url
 
@@ -15,11 +15,27 @@ from secator.utils import debug, escape_mongodb_url
 
 MONGODB_URL = CONFIG.addons.mongodb.url
 MONGODB_UPDATE_FREQUENCY = CONFIG.addons.mongodb.update_frequency
-MAX_POOL_SIZE = 100
+MONGODB_CONNECT_TIMEOUT = CONFIG.addons.mongodb.server_selection_timeout_ms
+MONGODB_MAX_POOL_SIZE = CONFIG.addons.mongodb.max_pool_size
 
 logger = logging.getLogger(__name__)
 
-client = pymongo.MongoClient(escape_mongodb_url(MONGODB_URL), maxPoolSize=MAX_POOL_SIZE)
+client = pymongo.MongoClient(
+	escape_mongodb_url(MONGODB_URL),
+	maxPoolSize=MONGODB_MAX_POOL_SIZE,
+	serverSelectionTimeoutMS=MONGODB_CONNECT_TIMEOUT
+)
+
+
+def get_runner_dbg(runner):
+	"""Runner debug object"""
+	return {
+		runner.unique_name: runner.status,
+		'type': runner.config.type,
+		'class': runner.__class__.__name__,
+		'caller': runner.config.name,
+		**runner.context
+	}
 
 
 def update_runner(self):
@@ -27,25 +43,19 @@ def update_runner(self):
 	type = self.config.type
 	collection = f'{type}s'
 	update = self.toDict()
-	debug_obj = {'type': 'runner', 'name': self.name, 'status': self.status}
 	chunk = update.get('chunk')
 	_id = self.context.get(f'{type}_chunk_id') if chunk else self.context.get(f'{type}_id')
-	debug('update', sub='hooks.mongodb', id=_id, obj=update, obj_after=True, level=4)
+	debug('to_update', sub='hooks.mongodb', id=_id, obj=get_runner_dbg(self), obj_after=True, obj_breaklines=False, verbose=True)  # noqa: E501
 	start_time = time.time()
 	if _id:
-		delta = start_time - self.last_updated if self.last_updated else MONGODB_UPDATE_FREQUENCY
-		if self.last_updated and delta < MONGODB_UPDATE_FREQUENCY and self.status == 'RUNNING':
-			debug(f'skipped ({delta:>.2f}s < {MONGODB_UPDATE_FREQUENCY}s)',
-				  sub='hooks.mongodb', id=_id, obj=debug_obj, obj_after=False, level=3)
-			return
 		db = client.main
 		start_time = time.time()
 		db[collection].update_one({'_id': ObjectId(_id)}, {'$set': update})
 		end_time = time.time()
 		elapsed = end_time - start_time
 		debug(
-			f'[dim gold4]updated in {elapsed:.4f}s[/]', sub='hooks.mongodb', id=_id, obj=debug_obj, obj_after=False, level=2)
-		self.last_updated = start_time
+			f'[dim gold4]updated in {elapsed:.4f}s[/]', sub='hooks.mongodb', id=_id, obj=get_runner_dbg(self), obj_after=False)  # noqa: E501
+		self.last_updated_db = start_time
 	else:  # sync update and save result to runner object
 		runner = db[collection].insert_one(update)
 		_id = str(runner.inserted_id)
@@ -55,13 +65,16 @@ def update_runner(self):
 			self.context[f'{type}_id'] = _id
 		end_time = time.time()
 		elapsed = end_time - start_time
-		debug(f'created in {elapsed:.4f}s', sub='hooks.mongodb', id=_id, obj=debug_obj, obj_after=False, level=2)
+		debug(f'in {elapsed:.4f}s', sub='hooks.mongodb', id=_id, obj=get_runner_dbg(self), obj_after=False)
 
 
 def update_finding(self, item):
+	if type(item) not in FINDING_TYPES:
+		return item
 	start_time = time.time()
 	db = client.main
 	update = item.toDict()
+	_type = item._type
 	_id = ObjectId(item._uuid) if ObjectId.is_valid(item._uuid) else None
 	if _id:
 		finding = db['findings'].update_one({'_id': _id}, {'$set': update})
@@ -72,7 +85,14 @@ def update_finding(self, item):
 		status = 'CREATED'
 	end_time = time.time()
 	elapsed = end_time - start_time
-	debug(f'in {elapsed:.4f}s', sub='hooks.mongodb', id=str(item._uuid), obj={'finding': status}, obj_after=False)
+	debug_obj = {
+		_type: status,
+		'type': 'finding',
+		'class': self.__class__.__name__,
+		'caller': self.config.name,
+		**self.context
+	}
+	debug(f'in {elapsed:.4f}s', sub='hooks.mongodb', id=str(item._uuid), obj=debug_obj, obj_after=False)  # noqa: E501
 	return item
 
 
@@ -80,20 +100,23 @@ def find_duplicates(self):
 	ws_id = self.toDict().get('context', {}).get('workspace_id')
 	if not ws_id:
 		return
-	celery_id = tag_duplicates.delay(ws_id)
-	debug(f'running duplicate check on workspace {ws_id}', id=celery_id, sub='hooks.mongodb')
+	if self.sync:
+		debug(f'running duplicate check on workspace {ws_id}', sub='hooks.mongodb')
+		tag_duplicates(ws_id)
+	else:
+		celery_id = tag_duplicates.delay(ws_id)
+		debug(f'running duplicate check on workspace {ws_id}', id=celery_id, sub='hooks.mongodb')
 
 
 def load_finding(obj):
 	finding_type = obj['_type']
 	klass = None
-	for otype in OUTPUT_TYPES:
+	for otype in FINDING_TYPES:
 		if finding_type == otype.get_name():
 			klass = otype
 			item = klass.load(obj)
 			item._uuid = str(obj['_id'])
 			return item
-	debug('could not load Secator output type from MongoDB object', obj=obj, sub='hooks.mongodb')
 	return None
 
 
@@ -149,18 +172,19 @@ def tag_duplicates(ws_id: str = None):
 				'seen dupes': len(seen_dupes)
 			},
 			id=ws_id,
-			sub='hooks.mongodb')
+			sub='hooks.mongodb.duplicates',
+			verbose=True)
 		tmp_duplicates_ids = list(dict.fromkeys([i._uuid for i in tmp_duplicates]))
-		debug(f'duplicate ids: {tmp_duplicates_ids}', id=ws_id, sub='hooks.mongodb')
+		debug(f'duplicate ids: {tmp_duplicates_ids}', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
 
 		# Update latest object as non-duplicate
 		if tmp_duplicates:
 			duplicates.extend([f for f in tmp_duplicates])
 			db.findings.update_one({'_id': ObjectId(item._uuid)}, {'$set': {'_related': tmp_duplicates_ids}})
-			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb')
+			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
 			non_duplicates.append(item)
 		else:
-			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb')
+			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
 			non_duplicates.append(item)
 
 	# debug(f'found {len(duplicates)} total duplicates')
@@ -184,19 +208,21 @@ def tag_duplicates(ws_id: str = None):
 			'duplicates': len(duplicates_ids),
 			'non-duplicates': len(non_duplicates_ids)
 		},
-		sub='hooks.mongodb')
+		sub='hooks.mongodb.duplicates')
 
 
-MONGODB_HOOKS = {
+HOOKS = {
 	Scan: {
+		'on_init': [update_runner],
 		'on_start': [update_runner],
-		'on_iter': [update_runner],
+		'on_interval': [update_runner],
 		'on_duplicate': [update_finding],
 		'on_end': [update_runner],
 	},
 	Workflow: {
+		'on_init': [update_runner],
 		'on_start': [update_runner],
-		'on_iter': [update_runner],
+		'on_interval': [update_runner],
 		'on_duplicate': [update_finding],
 		'on_end': [update_runner],
 	},
@@ -205,7 +231,7 @@ MONGODB_HOOKS = {
 		'on_start': [update_runner],
 		'on_item': [update_finding],
 		'on_duplicate': [update_finding],
-		'on_iter': [update_runner],
+		'on_interval': [update_runner],
 		'on_end': [update_runner, find_duplicates]
 	}
 }