feat(runner): add GCS driver and secator threads (#476)

Dockerfile

@@ -37,7 +37,8 @@ COPY . /code/
 RUN pipx install .
 RUN secator install tools
 RUN secator install addons worker
-RUN secator install addons google
+RUN secator install addons gdrive
+RUN secator install addons gcs
 RUN secator install addons mongodb
 RUN secator install addons redis
 RUN secator install addons dev

pyproject.toml

@@ -81,10 +81,13 @@ redis = [
 mongodb = [
   'pymongo < 5',
 ]
-google = [
+gdrive = [
   'google-api-python-client < 3',
   'gspread < 7'
 ]
+gcs = [
+	'google-cloud-storage < 3'
+]
 
 [project.scripts]
 secator = 'secator.cli:cli'

secator/cli.py

@@ -841,14 +841,13 @@ def health(json, debug):
 	console.print('\n:wrench: [bold gold3]Checking installed addons ...[/]')
 	table = get_health_table()
 	with Live(table, console=console):
-		for addon in ['worker', 'google', 'mongodb', 'redis', 'dev', 'trace', 'build']:
-			addon_var = ADDONS_ENABLED[addon]
+		for addon, installed in ADDONS_ENABLED.items():
 			info = {
 				'name': addon,
 				'version': None,
-				'status': 'ok' if addon_var else 'missing',
+				'status': 'ok' if installed else 'missing',
 				'latest_version': None,
-				'installed': addon_var,
+				'installed': installed,
 				'location': None
 			}
 			row = fmt_health_table_row(info, 'addons')
@@ -898,7 +897,7 @@ def run_install(cmd, title, next_steps=None):
 		if ret.return_code != 0:
 			console.print(f':exclamation_mark: Failed to install {title}.', style='bold red')
 		else:
-			console.print(f':tada: {title.capitalize()} installed successfully !', style='bold green')
+			console.print(f':tada: {title} installed successfully !', style='bold green')
 			if next_steps:
 				console.print('[bold gold3]:wrench: Next steps:[/]')
 				for ix, step in enumerate(next_steps):
@@ -920,10 +919,10 @@ def addons():
 
 @addons.command('worker')
 def install_worker():
-	"Install worker addon."
+	"Install Celery worker addon."
 	run_install(
 		cmd=f'{sys.executable} -m pip install secator[worker]',
-		title='worker addon',
+		title='Celery worker addon',
 		next_steps=[
 			'Run [bold green4]secator worker[/] to run a Celery worker using the file system as a backend and broker.',
 			'Run [bold green4]secator x httpx testphp.vulnweb.com[/] to admire your task running in a worker.',
@@ -932,26 +931,38 @@ def install_worker():
 	)
 
 
-@addons.command('google')
-def install_google():
-	"Install google addon."
+@addons.command('gdrive')
+def install_gdrive():
+	"Install Google Drive addon."
 	run_install(
 		cmd=f'{sys.executable} -m pip install secator[google]',
-		title='google addon',
+		title='Google Drive addon',
 		next_steps=[
-			'Run [bold green4]secator config set addons.google.credentials_path <VALUE>[/].',
-			'Run [bold green4]secator config set addons.google.drive_parent_folder_id <VALUE>[/].',
+			'Run [bold green4]secator config set addons.gdrive.credentials_path <VALUE>[/].',
+			'Run [bold green4]secator config set addons.gdrive.drive_parent_folder_id <VALUE>[/].',
 			'Run [bold green4]secator x httpx testphp.vulnweb.com -o gdrive[/] to send reports to Google Drive.'
 		]
 	)
 
 
+@addons.command('gcs')
+def install_gcs():
+	"Install Google Cloud Storage addon."
+	run_install(
+		cmd=f'{sys.executable} -m pip install secator[gcs]',
+		title='Google Cloud Storage addon',
+		next_steps=[
+			'Run [bold green4]secator config set addons.gcs.credentials_path <VALUE>[/].',
+		]
+	)
+
+
 @addons.command('mongodb')
 def install_mongodb():
-	"Install mongodb addon."
+	"Install MongoDB addon."
 	run_install(
 		cmd=f'{sys.executable} -m pip install secator[mongodb]',
-		title='mongodb addon',
+		title='MongoDB addon',
 		next_steps=[
 			'[dim]\[optional][/] Run [bold green4]docker run --name mongo -p 27017:27017 -d mongo:latest[/] to run a local MongoDB instance.',  # noqa: E501
 			'Run [bold green4]secator config set addons.mongodb.url mongodb://<URL>[/].',
@@ -962,10 +973,10 @@ def install_mongodb():
 
 @addons.command('redis')
 def install_redis():
-	"Install redis addon."
+	"Install Redis addon."
 	run_install(
 		cmd=f'{sys.executable} -m pip install secator[redis]',
-		title='redis addon',
+		title='Redis addon',
 		next_steps=[
 			'[dim]\[optional][/] Run [bold green4]docker run --name redis -p 6379:6379 -d redis[/] to run a local Redis instance.',  # noqa: E501
 			'Run [bold green4]secator config set celery.broker_url redis://<URL>[/]',

secator/config.py

@@ -121,12 +121,18 @@ class Wordlists(StrictModel):
 	lists: Dict[str, List[str]] = {}
 
 
-class GoogleAddon(StrictModel):
+class GoogleDriveAddon(StrictModel):
 	enabled: bool = False
 	drive_parent_folder_id: str = ''
 	credentials_path: str = ''
 
 
+class GoogleCloudStorageAddon(StrictModel):
+	enabled: bool = False
+	bucket_name: str = ''
+	credentials_path: str = ''
+
+
 class WorkerAddon(StrictModel):
 	enabled: bool = False
 
@@ -140,7 +146,8 @@ class MongodbAddon(StrictModel):
 
 
 class Addons(StrictModel):
-	google: GoogleAddon = GoogleAddon()
+	gdrive: GoogleDriveAddon = GoogleDriveAddon()
+	gcs: GoogleCloudStorageAddon = GoogleCloudStorageAddon()
 	worker: WorkerAddon = WorkerAddon()
 	mongodb: MongodbAddon = MongodbAddon()
 
@@ -170,7 +177,7 @@ class Config(DotMap):
 	>>> config = Config.parse(path='/path/to/config.yml')  # get custom config (from YAML file).
 	>>> config.print() 									   # print config without defaults.
 	>>> config.print(partial=False)  					   # print full config.
-	>>> config.set('addons.google.enabled', False)         # set value in config.
+	>>> config.set('addons.gdrive.enabled', False)         # set value in config.
 	>>> config.save()									   # save config back to disk.
 	"""
 

secator/decorators.py

@@ -317,8 +317,37 @@ def func(ctx, **opts):
 		# unknown_opts = get_unknown_opts(ctx)
 		# opts.update(unknown_opts)
 
+		# Expand input
 		inputs = opts.pop(input_type)
 		inputs = expand_input(inputs, ctx)
+
+		# Build hooks from driver name
+		hooks = []
+		drivers = driver.split(',') if driver else []
+		console = _get_rich_console()
+		supported_drivers = ['mongodb', 'gcs']
+		for driver in drivers:
+			if driver in supported_drivers:
+				if not ADDONS_ENABLED[driver]:
+					console.print(f'[bold red]Missing "{driver}" addon: please run `secator install addons {driver}`[/].')
+					sys.exit(1)
+				from secator.utils import import_dynamic
+				driver_hooks = import_dynamic(f'secator.hooks.{driver}', 'HOOKS')
+				if driver_hooks is None:
+					console.print(f'[bold red]Missing "secator.hooks.{driver}.HOOKS".[/]')
+					sys.exit(1)
+				hooks.append(driver_hooks)
+			else:
+				supported_drivers_str = ', '.join([f'[bold green]{_}[/]' for _ in supported_drivers])
+				console.print(f'[bold red]Driver "{driver}" is not supported.[/]')
+				console.print(f'Supported drivers: {supported_drivers_str}')
+				sys.exit(1)
+
+		from secator.utils import deep_merge_dicts
+		hooks = deep_merge_dicts(*hooks)
+		print(hooks)
+
+		# Enable sync or not
 		if sync or show:
 			sync = True
 		else:
@@ -335,15 +364,6 @@ def func(ctx, **opts):
 						_get_rich_console().print('[bold red]Missing `redis` addon: please run `secator install addons redis`[/].')
 						sys.exit(1)
 
-		# Build hooks from driver name
-		hooks = {}
-		if driver == 'mongodb':
-			if not ADDONS_ENABLED['mongodb']:
-				_get_rich_console().print('[bold red]Missing `mongodb` addon: please run `secator install addons mongodb`[/].')
-				sys.exit(1)
-			from secator.hooks.mongodb import MONGODB_HOOKS
-			hooks = MONGODB_HOOKS
-
 		# Set run options
 		opts.update({
 			'print_cmd': True,

secator/definitions.py

@@ -125,7 +125,8 @@ def is_importable(module_to_import):
 
 for addon, module in [
 	('worker', 'eventlet'),
-	('google', 'gspread'),
+	('gdrive', 'gspread'),
+	('gcs', 'google.cloud.storage'),
 	('mongodb', 'pymongo'),
 	('redis', 'redis'),
 	('dev', 'flake8'),

secator/exporters/gdrive.py

@@ -4,7 +4,7 @@
 
 from secator.config import CONFIG
 from secator.exporters._base import Exporter
-from secator.output_types import Info
+from secator.output_types import Info, Error
 from secator.rich import console
 from secator.utils import pluralize
 
@@ -17,20 +17,22 @@ def send(self):
 		title = self.report.data['info']['title']
 		sheet_title = f'{self.report.data["info"]["title"]}_{self.report.timestamp}'
 		results = self.report.data['results']
-		if not CONFIG.addons.google.credentials_path:
-			console.print(':file_cabinet: Missing CONFIG.addons.google.credentials_path to save to Google Sheets', style='red')
+		if not CONFIG.addons.gdrive.credentials_path:
+			error = Error('Missing CONFIG.addons.gdrive.credentials_path to save to Google Sheets')
+			console.print(error)
 			return
-		if not CONFIG.addons.google.drive_parent_folder_id:
-			console.print(':file_cabinet: Missing CONFIG.addons.google.drive_parent_folder_id to save to Google Sheets.', style='red')  # noqa: E501
+		if not CONFIG.addons.gdrive.drive_parent_folder_id:
+			error = Error('Missing CONFIG.addons.gdrive.drive_parent_folder_id to save to Google Sheets.')
+			console.print(error)
 			return
-		client = gspread.service_account(CONFIG.addons.google.credentials_path)
+		client = gspread.service_account(CONFIG.addons.gdrive.credentials_path)
 
 		# Create workspace folder if it doesn't exist
-		folder_id = self.get_folder_by_name(ws, parent_id=CONFIG.addons.google.drive_parent_folder_id)
+		folder_id = self.get_folder_by_name(ws, parent_id=CONFIG.addons.gdrive.drive_parent_folder_id)
 		if ws and not folder_id:
 			folder_id = self.create_folder(
 				folder_name=ws,
-				parent_id=CONFIG.addons.google.drive_parent_folder_id)
+				parent_id=CONFIG.addons.gdrive.drive_parent_folder_id)
 
 		# Create worksheet
 		sheet = client.create(title, folder_id=folder_id)
@@ -58,8 +60,9 @@ def send(self):
 			]
 			csv_path = f'{self.report.output_folder}/report_{output_type}.csv'
 			if not os.path.exists(csv_path):
-				console.print(
+				error = Error(
 					f'Unable to find CSV at {csv_path}. For Google sheets reports, please enable CSV reports as well.')
+				console.print(error)
 				return
 			sheet_title = pluralize(output_type).upper()
 			ws = sheet.add_worksheet(sheet_title, rows=len(items), cols=len(keys))
@@ -86,7 +89,7 @@ def send(self):
 	def create_folder(self, folder_name, parent_id=None):
 		from googleapiclient.discovery import build
 		from google.oauth2 import service_account
-		creds = service_account.Credentials.from_service_account_file(CONFIG.addons.google.credentials_path)
+		creds = service_account.Credentials.from_service_account_file(CONFIG.addons.gdrive.credentials_path)
 		service = build('drive', 'v3', credentials=creds)
 		body = {
 			'name': folder_name,
@@ -100,7 +103,7 @@ def create_folder(self, folder_name, parent_id=None):
 	def list_folders(self, parent_id):
 		from googleapiclient.discovery import build
 		from google.oauth2 import service_account
-		creds = service_account.Credentials.from_service_account_file(CONFIG.addons.google.credentials_path)
+		creds = service_account.Credentials.from_service_account_file(CONFIG.addons.gdrive.credentials_path)
 		service = build('drive', 'v3', credentials=creds)
 		driveid = service.files().get(fileId='root').execute()['id']
 		response = service.files().list(

secator/hooks/gcs.py

@@ -0,0 +1,49 @@
+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 Path(v).exists():
+			blob_name = f'{item._uuid}_{k}'
+			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

@@ -65,7 +65,7 @@ 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=get_runner_dbg(self), obj_after=False)
+		debug(f'in {elapsed:.4f}s', sub='hooks.mongodb', id=_id, obj=get_runner_dbg(self), obj_after=False)
 
 
 def update_finding(self, item):
@@ -211,7 +211,7 @@ def tag_duplicates(ws_id: str = None):
 		sub='hooks.mongodb')
 
 
-MONGODB_HOOKS = {
+HOOKS = {
 	Scan: {
 		'on_init': [update_runner],
 		'on_start': [update_runner],

secator/output_types/error.py

@@ -20,7 +20,13 @@ class Error(OutputType):
 	_sort_by = ('_timestamp',)
 
 	def from_exception(e, **kwargs):
-		return Error(message=f'{type(e).__name__}: {str(e)}', traceback=traceback_as_string(e), **kwargs)
+		message = type(e).__name__
+		if str(e):
+			message += f': {str(e)}'
+		return Error(message=message, traceback=traceback_as_string(e), **kwargs)
+
+	def __str__(self):
+		return self.message
 
 	def __repr__(self):
 		s = f'[bold red]❌ {self.message}[/]'