feat(runner): on serialized hooks (#424)

package.json

@@ -6,7 +6,7 @@
 	"scripts": {
 		"solo": "npm run venv && venv/bin/secator worker -r",
 		"dev": "npm run venv && SECATOR_CELERY_BROKER_URL=redis://localhost:6379 SECATOR_CELERY_RESULT_BACKEND=redis://localhost:6379 venv/bin/secator worker -r",
-		"venv": "pip install virtualenv --break-system-packages && virtualenv venv && chmod +x venv/bin/activate && . venv/bin/activate && venv/bin/pip install -e .[dev,worker,redis,mongodb,trace,dev]",
+		"venv": "pip install virtualenv --break-system-packages && virtualenv venv && chmod +x venv/bin/activate && . venv/bin/activate && venv/bin/pip install -e .[dev,worker,redis,mongodb,trace]",
 		"generate": "rm -r venv && npm run venv && venv/bin/pip install fastapi uvicorn && venv/bin/pip freeze > requirements.txt",
 		"docker:build": "docker build -t secator .",
 		"docker:push": "gcloud builds submit .",

secator/celery.py

@@ -145,8 +145,7 @@ def break_task(task_cls, task_opts, targets, results=[], chunk_size=1):
 
 @app.task(bind=True)
 def run_task(self, args=[], kwargs={}):
-	if CONFIG.debug.level > 1:
-		logger.info(f'Received task with args {args} and kwargs {kwargs}')
+	debug(f'Received task with args {args} and kwargs {kwargs}', sub="celery", level=2)
 	if 'context' not in kwargs:
 		kwargs['context'] = {}
 	kwargs['context']['celery_id'] = self.request.id
@@ -156,8 +155,7 @@ def run_task(self, args=[], kwargs={}):
 
 @app.task(bind=True)
 def run_workflow(self, args=[], kwargs={}):
-	if CONFIG.debug.level > 1:
-		logger.info(f'Received workflow with args {args} and kwargs {kwargs}')
+	debug(f'Received workflow with args {args} and kwargs {kwargs}', sub="celery", level=2)
 	if 'context' not in kwargs:
 		kwargs['context'] = {}
 	kwargs['context']['celery_id'] = self.request.id
@@ -167,8 +165,7 @@ def run_workflow(self, args=[], kwargs={}):
 
 @app.task(bind=True)
 def run_scan(self, args=[], kwargs={}):
-	if CONFIG.debug.level > 1:
-		logger.info(f'Received scan with args {args} and kwargs {kwargs}')
+	debug(f'Received scan with args {args} and kwargs {kwargs}', sub="celery", level=2)
 	if 'context' not in kwargs:
 		kwargs['context'] = {}
 	kwargs['context']['celery_id'] = self.request.id

secator/runners/command.py

@@ -422,22 +422,16 @@ def yielder(self):
 				line = self.run_hooks('on_line', line)
 
 				# Run item_loader to try parsing as dict
-				items = None
+				item_count = 0
 				if self.output_json:
-					items = self.run_item_loaders(line)
+					for item in self.run_item_loaders(line):
+						yield item
+						item_count += 1
 
-				# Yield line if no items parsed
-				if not items:
+				# Yield line if no items were yielded
+				if item_count == 0:
 					yield line
 
-				# Turn results into list if not already a list
-				elif not isinstance(items, list):
-					items = [items]
-
-				# Yield items
-				if items:
-					yield from items
-
 		except KeyboardInterrupt:
 			self.process.kill()
 			self.killed = True
@@ -446,19 +440,16 @@ def yielder(self):
 		self._wait_for_end()
 
 	def run_item_loaders(self, line):
-		"""Run item loaders on a string."""
-		items = []
+		"""Run item loaders against an output line."""
 		for item_loader in self.item_loaders:
-			result = None
 			if (callable(item_loader)):
-				result = item_loader(self, line)
+				yield from item_loader(self, line)
 			elif item_loader:
-				result = item_loader.run(line)
-			if isinstance(result, dict):
-				result = [result]
-			if result:
-				items.extend(result)
-		return items
+				name = item_loader.__class__.__name__.replace('Serializer', '').lower()
+				default_callback = lambda self, x: [(yield x)]  # noqa: E731
+				callback = getattr(self, f'on_{name}_loaded', None) or default_callback
+				for item in item_loader.run(line):
+					yield from callback(self, item)
 
 	def _prompt_sudo(self, command):
 		"""

secator/serializers/json.py

@@ -7,9 +7,9 @@ def run(self, line):
 		start_index = line.find('{')
 		end_index = line.rfind('}')
 		if start_index == -1 or end_index == -1:
-			return None
+			return
 		try:
 			json_obj = line[start_index:end_index+1]
-			return yaml.safe_load(json_obj)
+			yield yaml.safe_load(json_obj)
 		except yaml.YAMLError:
-			return None
+			return

secator/serializers/regex.py

@@ -11,7 +11,7 @@ def run(self, line):
 		match = self.regex.match(line)
 		output = {}
 		if not match:
-			return None
+			return
 		for field in self.fields:
 			output[field] = match.group(field)
-		return output
+		yield output

secator/tasks/cariddi.py

@@ -1,5 +1,3 @@
-import json
-
 from secator.decorators import task
 from secator.definitions import (DELAY, DEPTH, FILTER_CODES, FILTER_REGEX,
 							   FILTER_SIZE, FILTER_WORDS, FOLLOW_REDIRECT,
@@ -41,7 +39,6 @@ class cariddi(HttpCrawler):
 		TIMEOUT: 't',
 		USER_AGENT: 'ua'
 	}
-	item_loaders = []
 	install_cmd = 'go install -v github.com/edoardottt/cariddi/cmd/cariddi@latest'
 	install_github_handle = 'edoardottt/cariddi'
 	encoding = 'ansi'
@@ -51,53 +48,45 @@ class cariddi(HttpCrawler):
 	profile = 'cpu'
 
 	@staticmethod
-	def item_loader(self, line):
-		items = []
-		try:
-			item = json.loads(line)
-			url_item = {k: v for k, v in item.items() if k != 'matches'}
-			url = url_item[URL]
-			items.append(url_item)
-			matches = item.get('matches', {})
-			params = matches.get('parameters', [])
-			errors = matches.get('errors', [])
-			secrets = matches.get('secrets', [])
-			infos = matches.get('infos', [])
-
-			for param in params:
-				param_name = param['name']
-				for attack in param['attacks']:
-					extra_data = {'param': param_name, 'source': 'url'}
-					item = {
-						'name': attack + ' param',
-						'match': url,
-						'extra_data': extra_data
-					}
-					items.append(item)
-
-			for error in errors:
-				match = error['match']
-				match = (match[:1000] + '...TRUNCATED') if len(match) > 1000 else match  # truncate as this can be a very long match
-				error['extra_data'] = {'error': match, 'source': 'body'}
-				error['match'] = url
-				items.append(error)
+	def on_json_loaded(self, item):
+		url_item = {k: v for k, v in item.items() if k != 'matches'}
+		url = url_item[URL]
+		yield url_item
+		matches = item.get('matches', {})
+		params = matches.get('parameters', [])
+		errors = matches.get('errors', [])
+		secrets = matches.get('secrets', [])
+		infos = matches.get('infos', [])
 
-			for secret in secrets:
-				match = secret['match']
-				secret['extra_data'] = {'secret': match, 'source': 'body'}
-				secret['match'] = url
-				items.append(secret)
+		for param in params:
+			param_name = param['name']
+			for attack in param['attacks']:
+				extra_data = {'param': param_name, 'source': 'url'}
+				parameter = {
+					'name': attack + ' param',
+					'match': url,
+					'extra_data': extra_data
+				}
+				yield parameter
 
-			for info in infos:
-				CARIDDI_IGNORE_LIST = ['BTC address']  # TODO: make this a config option
-				if info['name'] in CARIDDI_IGNORE_LIST:
-					continue
-				match = info['match']
-				info['extra_data'] = {'info': match, 'source': 'body'}
-				info['match'] = url
-				items.append(info)
+		for error in errors:
+			match = error['match']
+			match = (match[:1000] + '...TRUNCATED') if len(match) > 1000 else match  # truncate as this can be a very long match
+			error['extra_data'] = {'error': match, 'source': 'body'}
+			error['match'] = url
+			yield error
 
-		except json.decoder.JSONDecodeError:
-			pass
+		for secret in secrets:
+			match = secret['match']
+			secret['extra_data'] = {'secret': match, 'source': 'body'}
+			secret['match'] = url
+			yield secret
 
-		return items
+		for info in infos:
+			CARIDDI_IGNORE_LIST = ['BTC address']  # TODO: make this a config option
+			if info['name'] in CARIDDI_IGNORE_LIST:
+				continue
+			match = info['match']
+			info['extra_data'] = {'info': match, 'source': 'body'}
+			info['match'] = url
+			yield info

secator/tasks/dnsx.py

@@ -2,7 +2,6 @@
 from secator.definitions import (OPT_PIPE_INPUT, RATE_LIMIT, RETRIES, THREADS)
 from secator.output_types import Record, Ip
 from secator.tasks._categories import ReconDns
-import json
 
 
 @task()
@@ -23,20 +22,12 @@ class dnsx(ReconDns):
 		'resolver': {'type': str, 'short': 'r', 'help': 'List of resolvers to use (file or comma separated)'},
 		'wildcard_domain': {'type': str, 'short': 'wd', 'help': 'Domain name for wildcard filtering'},
 	}
-	item_loaders = []
 	install_cmd = 'go install -v github.com/projectdiscovery/dnsx/cmd/dnsx@latest'
 	install_github_handle = 'projectdiscovery/dnsx'
 	profile = 'io'
 
 	@staticmethod
-	def item_loader(self, line):
-		try:
-			item = json.loads(line)
-		except json.JSONDecodeError:
-			return
-		if self.orig:  # original dnsx JSON output
-			yield item
-			return
+	def on_json_loaded(self, item):
 		host = item['host']
 		record_types = ['a', 'aaaa', 'cname', 'mx', 'ns', 'txt', 'srv', 'ptr', 'soa', 'axfr', 'caa']
 		for _type in record_types:

secator/tasks/fping.py

@@ -29,13 +29,14 @@ class fping(ReconIp):
 	}
 	input_type = IP
 	output_types = [Ip]
+	item_loaders = []
 	install_cmd = 'sudo apt install -y fping'
 
 	@staticmethod
 	def item_loader(self, line):
-		if validators.ipv4(line) or validators.ipv6(line):
-			return {'ip': line, 'alive': True}
-		return None
+		if not (validators.ipv4(line) or validators.ipv6(line)):
+			return
+		yield {'ip': line, 'alive': True}
 
 	@staticmethod
 	def on_line(self, line):

secator/tasks/gf.py

@@ -23,10 +23,11 @@ class gf(Tagger):
 		'git clone https://github.com/1ndianl33t/Gf-Patterns $HOME/.gf || true'
 	)
 	output_types = [Tag]
+	item_loaders = []
 
 	@staticmethod
 	def item_loader(self, line):
-		return {'match': line, 'name': self.get_opt_value('pattern').rstrip() + ' pattern'}  # noqa: E731,E501
+		yield {'match': line, 'name': self.get_opt_value('pattern').rstrip() + ' pattern'}  # noqa: E731,E501
 
 	@staticmethod
 	def on_item(self, item):

secator/tasks/grype.py

@@ -27,6 +27,7 @@ class grype(VulnCode):
 		USER_AGENT: OPT_NOT_SUPPORTED
 	}
 	output_types = [Vulnerability]
+	item_loaders = []
 	install_cmd = (
 		'curl -sSfL https://raw.githubusercontent.com/anchore/grype/main/install.sh | sudo sh -s -- -b /usr/local/bin'
 	)
@@ -37,7 +38,7 @@ def item_loader(self, line):
 		"""Load vulnerabilty dicts from grype line output."""
 		split = [i for i in line.split(' ') if i]
 		if not len(split) in [5, 6] or split[0] == 'NAME':
-			return None
+			return
 		version_fixed = None
 		if len(split) == 5:  # no version fixed
 			product, version, product_type, vuln_id, severity = tuple(split)
@@ -76,4 +77,4 @@ def item_loader(self, line):
 				data.update(vuln)
 				data['severity'] = data['severity'] or severity.lower()
 		data['extra_data'] = extra_data
-		return data
+		yield data

secator/tasks/katana.py

@@ -1,5 +1,4 @@
 import os
-import json
 from urllib.parse import urlparse
 
 from secator.decorators import task
@@ -74,7 +73,6 @@ class katana(HttpCrawler):
 			# TAGS: lambda x: x['response'].get('server')
 		}
 	}
-	item_loaders = []
 	install_cmd = 'sudo apt install build-essential && go install -v github.com/projectdiscovery/katana/cmd/katana@latest'
 	install_github_handle = 'projectdiscovery/katana'
 	proxychains = False
@@ -83,12 +81,7 @@ class katana(HttpCrawler):
 	profile = 'io'
 
 	@staticmethod
-	def item_loader(self, item):
-		try:
-			item = json.loads(item)
-		except json.JSONDecodeError:
-			return None
-
+	def on_json_loaded(self, item):
 		# form detection
 		forms = item.get('response', {}).get('forms', [])
 		if forms:

secator/tasks/mapcidr.py

@@ -17,6 +17,7 @@ class mapcidr(ReconIp):
 	install_github_handle = 'projectdiscovery/mapcidr'
 	input_type = CIDR_RANGE
 	output_types = [Ip]
+	item_loaders = []
 	opt_key_map = {
 		THREADS: OPT_NOT_SUPPORTED,
 		PROXY: OPT_NOT_SUPPORTED,
@@ -29,5 +30,5 @@ class mapcidr(ReconIp):
 	@staticmethod
 	def item_loader(self, line):
 		if validators.ipv4(line) or validators.ipv6(line):
-			return {'ip': line, 'alive': False}
-		return None
+			yield {'ip': line, 'alive': False}
+		return