refactor: integration tests

secator/celery.py

@@ -251,7 +251,7 @@ def run_command(self, results, name, targets, opts={}):
 	# Get expanded targets
 	if not chunk and results:
 		targets, opts = run_extractors(results, opts, targets)
-		debug('after extractors', obj={'targets': targets, 'opts': opts}, sub='celery.state')
+		debug('after extractors', obj={'targets': targets}, sub='celery.inputs', verbose=True)
 
 	try:
 		# Get task class

secator/configs/workflows/url_fuzz.yaml

@@ -23,6 +23,7 @@ tasks:
     targets_:
       type: url
       field: url
+      condition: "item.status_code != 0"
   katana:
     description: Run crawler on found directories
     targets_:

secator/output_types/error.py

@@ -23,7 +23,8 @@ def from_exception(e, **kwargs):
 		message = type(e).__name__
 		if str(e):
 			message += f': {str(e)}'
-		return Error(message=message, traceback=traceback_as_string(e), **kwargs)
+		tb = not isinstance(e, KeyboardInterrupt)
+		return Error(message=message, traceback=traceback_as_string(e) if tb else '', **kwargs)
 
 	def __str__(self):
 		return self.message

secator/runners/_base.py

@@ -557,7 +557,7 @@ def run_validators(self, validator_type, *args, error=True):
 			name = f'{self.__class__.__name__}.{validator_type}'
 			fun = self.get_func_path(validator)
 			if not validator(self, *args):
-				self.debug('', obj={name + ' [dim yellow]->[/] ' + fun: 'failed'}, id=_id, sub='validators')
+				self.debug('', obj={name + ' [dim yellow]->[/] ' + fun: 'failed'}, id=_id, sub='validators', verbose=True)
 				doc = validator.__doc__
 				if error:
 					message = 'Validator failed'
@@ -570,7 +570,7 @@ def run_validators(self, validator_type, *args, error=True):
 					)
 					self.add_result(error, print=True)
 				return False
-			self.debug('', obj={name + ' [dim yellow]->[/] ' + fun: 'success'}, id=_id, sub='validators')
+			self.debug('', obj={name + ' [dim yellow]->[/] ' + fun: 'success'}, id=_id, sub='validators', verbose=True)
 		return True
 
 	def register_hooks(self, hooks):
@@ -793,8 +793,8 @@ def _process_item(self, item):
 			return
 
 		# Abort further processing if no_process is set
-		if self.no_process:
-			return
+		# if self.no_process:
+		# 	return
 
 		# Run item validators
 		if not self.run_validators('validate_item', item, error=False):

secator/runners/command.py

@@ -361,7 +361,6 @@ def yielder(self):
 
 			# Output and results
 			self.return_code = 0
-			self.killed = False
 
 			# Run the command using subprocess
 			env = os.environ
@@ -397,9 +396,9 @@ def yielder(self):
 			yield from self.handle_file_not_found(e)
 
 		except BaseException as e:
-			self.debug(f'{self.unique_name}: {type(e).__name__}.', sub='error')
-			self.stop_process()
+			self.debug(f'{type(e).__name__}', sub='error')
 			yield Error.from_exception(e, _source=self.unique_name, _uuid=str(uuid.uuid4()))
+			self.stop_process()
 
 		finally:
 			yield from self._wait_for_end()
@@ -598,26 +597,17 @@ def _wait_for_end(self):
 		"""Wait for process to finish and process output and return code."""
 		if not self.process:
 			return
+		self.debug('Waiting for command to finish ...', sub='command')
 		for line in self.process.stdout.readlines():
 			yield from self.process_line(line)
 		self.process.wait()
 		self.return_code = self.process.returncode
 		self.process.stdout.close()
 		self.return_code = 0 if self.ignore_return_code else self.return_code
 		self.output = self.output.strip()
-		self.killed = self.return_code == -2 or self.killed
-		self.debug(f'Command {self.cmd} finished with return code {self.return_code}', sub='command')
-
-		if self.killed:
-			error = 'Process was killed manually (CTRL+C / CTRL+X)'
-			yield Error(
-				message=error,
-				_source=self.unique_name,
-				_uuid=str(uuid.uuid4())
-			)
-
-		elif self.return_code != 0:
-			error = f'Command failed with return code {self.return_code}.'
+		self.debug(f'finished with return code {self.return_code}', sub='command')
+		if self.return_code != 0 and self.return_code != 130:
+			error = f'failed with return code {self.return_code}.'
 			last_lines = self.output.split('\n')
 			last_lines = last_lines[max(0, len(last_lines) - 2):]
 			yield Error(

secator/utils_test.py

@@ -4,6 +4,9 @@
 import sys
 import unittest.mock
 
+from threading import Thread
+from time import sleep
+
 from fp.fp import FreeProxy
 
 from secator.definitions import (CIDR_RANGE, DELAY, DEPTH, EMAIL,
@@ -22,6 +25,7 @@
 USE_PROXY = bool(int(os.environ.get('USE_PROXY', '0')))
 TEST_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + '/tests/'
 FIXTURES_DIR = f'{TEST_DIR}/fixtures'
+INTEGRATION_DIR = f'{TEST_DIR}/integration'
 USE_PROXY = bool(int(os.environ.get('USE_PROXY', '0')))
 
 #------------#
@@ -210,6 +214,48 @@ def _test_runner_output(
 		console.print('[dim green] ok[/]')
 
 
+class SecatorTestCase(unittest.TestCase):
+
+	integration_lab = True
+	celery_worker = False
+	mongodb = False
+	redis = False
+
+	@classmethod
+	def setUpClass(cls):
+		cls.threads = []
+		cls.commands = []
+		commands = {
+			'integration_lab': 'docker-compose up',
+			'celery_worker': 'secator worker',
+			'mongodb': 'docker run --name integration-mongo -p 27018:27017 mongo:latest',
+			'redis': 'docker run --name integration-redis -p 6380:6379 redis'
+		}
+		for addon, cmd in commands.items():
+			if getattr(cls, addon, False) is True:
+				command = Command.execute(cmd, name=addon, cwd=INTEGRATION_DIR, quiet=False, run=False)
+				cls.commands.append(command)
+				cls.threads.append(Thread(target=command.run))
+
+		# Start all threads
+		[thread.start() for thread in cls.threads]
+
+		# Wait a bit and check for errors
+		sleep(20)
+		for command in cls.commands:
+			if command.status == 'FAILURE':
+				cls.tearDownClass()
+				assert False, f'{command.name} failed to start properly'
+
+	@classmethod
+	def tearDownClass(cls):
+		# Stop all command
+		[command.stop_process() for command in cls.commands]
+
+		# Stop all threads
+		[thread.join() for thread in cls.threads]
+
+
 def clear_modules():
 	"""Clear all secator modules imports.
 	See https://stackoverflow.com/questions/7460363/re-import-module-under-test-to-lose-context for context.

tests/integration/setup.sh

@@ -2,4 +2,4 @@
 
 echo "Setting up Juice shop ..."
 docker-compose pull
-docker-compose up -d
+docker-compose up

tests/integration/test_addons.py

@@ -1,23 +1,17 @@
 import os
-import unittest
 from unittest import mock
-from secator.utils_test import clear_modules
+from secator.utils_test import clear_modules, SecatorTestCase
 from secator.config import CONFIG
 
 
-class TestAddonMongo(unittest.TestCase):
+class TestAddonMongo(SecatorTestCase):
 
 	@classmethod
-	@mock.patch.dict(os.environ, {"SECATOR_ADDONS_MONGODB_URL": "mongodb://localhost"})
+	@mock.patch.dict(os.environ, {"SECATOR_ADDONS_MONGODB_URL": "mongodb://localhost:27018"})
 	def setUpClass(cls):
 		clear_modules()
 		from secator.config import CONFIG
-		print(CONFIG.addons.mongodb.url)
-		raise Exception('test')
-
-	@classmethod
-	def tearDownClass(cls):
-		pass
+		super().setUpClass()
 
 	def test_ok(self):
 		print(CONFIG.addons.mongodb.url)

tests/integration/test_celery.py

@@ -1,17 +1,10 @@
-import queue
 import os
-import unittest
-import warnings
-
-from time import sleep
-from threading import Thread
 
 from celery import chain, chord
 
 from secator.celery import app, forward_results
 from secator.config import CONFIG
-from secator.utils_test import TEST_TASKS, load_fixture
-from secator.runners import Command
+from secator.utils_test import SecatorTestCase, TEST_TASKS, load_fixture
 from secator.output_types import Url
 from tests.integration.inputs import INPUTS_SCANS
 
@@ -28,32 +21,9 @@
 HOST_TARGETS = INPUTS_SCANS['host']
 
 
-class TestCelery(unittest.TestCase):
+class TestCelery(SecatorTestCase):
 
-	@classmethod
-	def setUpClass(cls):
-		warnings.simplefilter('ignore', category=ResourceWarning)
-		warnings.simplefilter('ignore', category=DeprecationWarning)
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/setup.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
-		cls.queue = queue.Queue()
-		cls.cmd = Command.execute('secator worker', quiet=True, run=False)
-		cls.thread = Thread(target=cls.cmd.run)
-		cls.thread.start()
-		sleep(5)
-
-	@classmethod
-	def tearDownClass(cls) -> None:
-		cls.cmd.stop_process()
-		cls.thread.join()
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/teardown.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
+	celery_worker = True
 
 	def test_httpx_chain(self):
 		from secator.tasks import httpx

tests/integration/test_scans.py

@@ -1,14 +1,10 @@
 import logging
 import os
-import unittest
-import warnings
-from time import sleep
 
 from secator.definitions import DEBUG
-from secator.rich import console
-from secator.runners import Command, Scan
-from secator.utils import setup_logging, merge_opts
-from secator.utils_test import TEST_SCANS, CommandOutputTester, load_fixture
+from secator.runners import Scan
+from secator.utils import setup_logging
+from secator.utils_test import TEST_SCANS, CommandOutputTester, load_fixture, SecatorTestCase
 from tests.integration.inputs import INPUTS_SCANS
 from tests.integration.outputs import OUTPUTS_SCANS
 
@@ -17,24 +13,9 @@
 setup_logging(level)
 
 
-class TestScans(unittest.TestCase, CommandOutputTester):
+class TestScans(SecatorTestCase, CommandOutputTester):
 
-	def setUp(self):
-		warnings.simplefilter('ignore', category=ResourceWarning)
-		warnings.simplefilter('ignore', category=DeprecationWarning)
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/setup.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
-		sleep(15)
-
-	def tearDown(self):
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/teardown.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
+	celery_worker = True
 
 	def test_scans(self):
 		opts = {
@@ -58,4 +39,4 @@ def test_scans(self):
 				scan = Scan(conf, inputs=inputs, run_opts=opts)
 				self._test_runner_output(
 					scan,
-					expected_results=outputs)
+					expected_results=outputs)

tests/integration/test_tasks.py

@@ -1,36 +1,14 @@
 import os
-import unittest
-import warnings
-from time import sleep
 
-from secator.rich import console
-from secator.runners import Command
 from secator.utils import merge_opts
-from secator.utils_test import (META_OPTS, TEST_TASKS, CommandOutputTester,
-                              load_fixture)
+from secator.utils_test import (META_OPTS, TEST_TASKS, CommandOutputTester, SecatorTestCase, load_fixture)
 from tests.integration.inputs import INPUTS_TASKS
 from tests.integration.outputs import OUTPUTS_TASKS
 
 INTEGRATION_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
-class TestTasks(unittest.TestCase, CommandOutputTester):
-	def setUp(self):
-		warnings.simplefilter('ignore', category=ResourceWarning)
-		warnings.simplefilter('ignore', category=DeprecationWarning)
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/setup.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
-		sleep(15)
-
-	def tearDown(self):
-		Command.execute(
-			f'sh {INTEGRATION_DIR}/teardown.sh',
-			quiet=True,
-			cwd=INTEGRATION_DIR
-		)
+class TestTasks(SecatorTestCase, CommandOutputTester):
 
 	def test_tasks(self):
 		opts = META_OPTS.copy()

tests/integration/test_worker.py

@@ -1,25 +1,16 @@
-import unittest
+from time import sleep
+from threading import Thread
+
 from secator.output_types import Url, Target, Port, Vulnerability, Info
 from secator.runners import Command
 from secator.serializers import JSONSerializer
-from time import sleep
-from threading import Thread
-import queue
+from secator.utils_test import SecatorTestCase
 
-class TestWorker(unittest.TestCase):
 
-	@classmethod
-	def setUpClass(cls):
-		cls.queue = queue.Queue()
-		cls.cmd = Command.execute('secator worker', name='secator_worker', quiet=True, run=False)
-		cls.thread = Thread(target=cls.cmd.run)
-		cls.thread.start()
-		sleep(3)
+class TestWorker(SecatorTestCase):
 
-	@classmethod
-	def tearDownClass(cls) -> None:
-		cls.cmd.stop_process()
-		cls.thread.join()
+	celery_worker = True
+	lab = False
 
 	def test_httpx(self):
 		cmd = Command.execute(