From 9d34c13a90f4f4870bcaba8bb2d496277f697b2c Mon Sep 17 00:00:00 2001 From: Italo Cunha Date: Sun, 18 Jun 2023 15:25:29 -0300 Subject: [PATCH] refactor tests --- scout/scout.py | 50 +++++---- scout/tests/testrun.py | 228 ++++++++++++++++++++--------------------- 2 files changed, 139 insertions(+), 139 deletions(-) diff --git a/scout/scout.py b/scout/scout.py index c90b3c5..f47d2b4 100644 --- a/scout/scout.py +++ b/scout/scout.py @@ -52,8 +52,7 @@ class Task(Protocol): @dataclasses.dataclass(frozen=True) class ScoutTask: label: str - commands: str | None = None - volumes: dict | None = None + command: tuple[str] | None = None aws_api_key: str | None = None aws_api_secret: str | None = None role_arn: str | None = None @@ -94,18 +93,37 @@ def get_docker_client() -> docker.DockerClient: def enqueue(self, taskcfg: Task) -> None: assert isinstance(taskcfg, ScoutTask) - outfp = self.config.output_dir / taskcfg.label os.makedirs(outfp, exist_ok=True) try: + command = taskcfg.command if taskcfg.command is not None else ( + "scout", + "aws", + "--no-browser", + "--result-format", + "json", + "--report-dir", + f"{OUTDIR_CONTAINER_MOUNT}", + "--logfile", + f"{OUTDIR_CONTAINER_MOUNT}/scout.log", + ) ctx = self.docker.containers.run( self.config.docker_image, - command=taskcfg.commands, + command=command, detach=True, labels={SCOUT_TASK_LABEL_KEY: taskcfg.label}, stdout=True, stderr=True, - volumes=taskcfg.volumes, + volumes={ + str(self.config.credentials_file): { + "bind": "/root/.aws/credentials", + "mode": "ro", + }, + str(outfp): { + "bind": OUTDIR_CONTAINER_MOUNT, + "mode": "rw", + }, + }, working_dir="/root", ) except docker.errors.APIError as e: @@ -118,16 +136,15 @@ def enqueue(self, taskcfg: Task) -> None: def shutdown(self, wait: bool = True) -> None: logging.info("Scout shutting down (wait=%s)", wait) self.running = False + self.handle_finished_containers() if not wait: with self.lock: for ctx, cfg in self.containers: logging.warning("Force-closing container for task %s", cfg.label) ctx.remove(force=True) - self.containers.clear() - else: - self.handle_finished_containers() - self.thread.join() - + self.containers.clear() + logging.info("Joining Scout polling thread") + self.thread.join() logging.info("Joined Scout polling thread, module shut down") def scout_polling_thread(self) -> None: @@ -140,17 +157,10 @@ def handle_finished_containers(self) -> None: completed = set() with self.lock: for ctx, cfg in self.containers: - try: - ctx.reload() - except docker.errors.NotFound: - logging.warning("Container not found: %s", cfg.label) - continue - + ctx.reload() if not ContainerState(ctx.status).is_done(): continue - assert cfg.label == ctx.labels[SCOUT_TASK_LABEL_KEY] - r = ctx.wait(timeout=self.config.docker_timeout) r["stdout"] = ctx.logs( stdout=True, stderr=False, timestamps=True @@ -158,7 +168,6 @@ def handle_finished_containers(self) -> None: r["stderr"] = ctx.logs( stdout=False, stderr=True, timestamps=True ).decode("utf8") - outfp = self.config.output_dir / cfg.label with gzip.open(outfp / "result.json.gz", "wt", encoding="utf8") as fd: json.dump(r, fd) @@ -169,14 +178,11 @@ def handle_finished_containers(self) -> None: r[DOCKER_STATUSCODE_KEY], ) completed.add((ctx, cfg)) - self.containers -= completed - logging.info( "Running %d ScoutSuite containers, waiting %d seconds to refresh", len(self.containers), self.config.docker_poll_interval, ) - for ctx, _cfg in completed: ctx.remove() diff --git a/scout/tests/testrun.py b/scout/tests/testrun.py index d00653e..c52be40 100755 --- a/scout/tests/testrun.py +++ b/scout/tests/testrun.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 +import argparse import logging import os import pathlib @@ -8,155 +9,148 @@ import scout -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument( - "-t1", "--test1", type=int, help="Test 1: Container removed after n seconds." -) -parser.add_argument( - "-t2", - "--test2", - type=int, - help="Test 2: Container with infinite execution removed after n seconds.", -) -parser.add_argument( - "-t3", - "--test3", - type=int, - help="Test 3: Container removed after returning an error.", -) - -parser.add_argument( - "-scout", - "--scout_suite", - action="store_true", - help="SCOUT: Running the Scout module.", -) -args = parser.parse_args() +ALPINE_IMAGE = "alpine" +NUM_CONTAINERS = 5 +BASE_SLEEP_DURATION = 8 +SLEEP_INCREMENT = 3 -CWD = pathlib.Path(os.getcwd()) -assert pathlib.Path(CWD / __file__).exists(), "Run from inside tests/ with ./testrun.py" -SCOUT_MOD_PATH = CWD.parent - - -DOCKER_STATUSCODE_KEY = "StatusCode" -OUTDIR_CONTAINER_MOUNT = "/root/output" -SCOUT_TASK_LABEL_KEY = "scout-task-id" +def create_parser(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--cred-file", + dest="cred_file", + metavar="FILE", + type=pathlib.Path, + help="Path to AWS credentials file [%(default)s]", + default=pathlib.Path("~/.aws/credentials").expanduser(), + required=False, + ) + parser.add_argument( + "-t", "--test", + action="append", + dest="tests", + type=int, + choices=[1, 2, 3], + help="The tests to run, can be used multiple times %(default)s", + default=[], + required=False, + ) + parser.add_argument( + "--run-scout", + action="store_true", + help="Run Scout on AWS account in addition to tests [%(default)s]", + default=False, + required=False, + ) + parser.add_argument( + "--outdir", + dest="outdir", + metavar="DIR", + type=pathlib.Path, + help="Where to store test output", + default=pathlib.Path("./test-output").absolute(), + required=False, + ) + return parser def callback(label: str, success: bool) -> None: logging.info("Callback for %s running, success=%s", label, success) -def test_scout(s): - taskcfg = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - f"scout aws --no-browser --result-format json \ - --report-dir {OUTDIR_CONTAINER_MOUNT} --logfile \ - {OUTDIR_CONTAINER_MOUNT}/scout.log", - volumes={ - "./dev/aws-credentials": {"bind": "/root/.aws/credentials", "mode": "ro"}, - "./data": {"bind": "/root/output", "mode": "rw"}, - }, - ) +def run_scout(cfg: scout.ScoutConfig): + logging.info("Starting Scout module") + sm = scout.Scout(cfg, callback) - logging.info("RUNNING TEST SCOUT") - s.enqueue(taskcfg) - logging.info("Task submitted") - s.shutdown() + taskcfg = scout.ScoutTask("scout-run") + logging.info("Running Scout") + sm.enqueue(taskcfg) + logging.info("Waiting for Scout to terminate") + sm.shutdown(wait=True) -def task_1(s, seconds: int): - taskcfg = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - f"sleep {seconds}", - ) +def run_test_1(cfg: scout.ScoutConfig) -> None: + logging.info("Starting Scout module") + sm = scout.Scout(cfg, callback) - taskcfg2 = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - f"sleep {seconds-2}", - ) + logging.info("Running test 1") + logging.info("Will start %d containers", NUM_CONTAINERS) + for i in range(1, NUM_CONTAINERS+1): + taskcfg = scout.ScoutTask( + f"test-1-{i}", + ("sleep", f"{BASE_SLEEP_DURATION + SLEEP_INCREMENT*i}"), + ) + sm.enqueue(taskcfg) - taskcfg3 = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - f"sleep {seconds-4}", - ) + logging.info("Waiting for all containers to terminate cleanly") + sm.shutdown(wait=True) + logging.info("Test 1 completed") - logging.info( - "RUNNING TEST 1 - The container was removed after n seconds, along with another set of tasks." - ) - s.enqueue(taskcfg) - s.enqueue(taskcfg2) - s.enqueue(taskcfg2) - s.enqueue(taskcfg3) - s.enqueue(taskcfg3) - logging.info("Task submitted") - s.shutdown() - - -def task_2(s, seconds: int): - taskcfg = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - "tail -f /dev/null", - ) - logging.info( - "RUNNING TEST 2 - Container with infinite execution removed after n seconds." - ) - s.enqueue(taskcfg) +def run_test_2(cfg: scout.ScoutConfig): + logging.info("Starting Scout module") + sm = scout.Scout(cfg, callback) - time.sleep(seconds) - logging.info("Task submitted") - s.shutdown(False) + logging.info("Running test 2") + logging.info("Will start %d containers", NUM_CONTAINERS) + for i in range(1, NUM_CONTAINERS+1): + taskcfg = scout.ScoutTask( + f"test-2-{i}", + ("sleep", f"{BASE_SLEEP_DURATION + SLEEP_INCREMENT*i}"), + ) + sm.enqueue(taskcfg) + logging.info("Sleeping %d seconds", BASE_SLEEP_DURATION) + time.sleep(BASE_SLEEP_DURATION) + logging.info("Force-quitting remaining containers") + sm.shutdown(wait=False) + logging.info("Test 2 completed") -def task_3(s, seconds: int): - taskcfg = scout.ScoutTask( - time.strftime("scout-%Y%m%d-%H%M%S"), - f"sh -c 'sleep {seconds} && false'", - ) - logging.info("RUNNING TEST 3 - Container removed after returning an error.") - s.enqueue(taskcfg) +def run_test_3(cfg: scout.ScoutConfig) -> None: + logging.info("Starting Scout module") + sm = scout.Scout(cfg, callback) + + logging.info("Running test 3") + logging.info("Will start %d containers", NUM_CONTAINERS) + for i in range(1, NUM_CONTAINERS+1): + taskcfg = scout.ScoutTask( + f"test-3-{i}", + ("sh", "-c", f"sleep {BASE_SLEEP_DURATION + SLEEP_INCREMENT*i} && false"), + ) + sm.enqueue(taskcfg) - logging.info("Task submitted") - s.shutdown() + logging.info("Waiting for all containers to fail") + sm.shutdown(wait=True) + logging.info("Test 3 completed") -scout_image = "rossja/ncc-scoutsuite:aws-latest" -generic_image = "alpine" +TEST_FUNCTIONS = [run_test_1, run_test_2, run_test_3] def main(): - logging.basicConfig(level=logging.INFO) - logging.info("Starting Scout module") + parser = create_parser() + args = parser.parse_args() + + logging.basicConfig(level=logging.DEBUG) + cfg = scout.ScoutConfig( - pathlib.Path(SCOUT_MOD_PATH / "dev/aws-credentials"), - pathlib.Path(SCOUT_MOD_PATH / "data"), - docker_image=scout_image, + args.cred_file, + args.outdir, + docker_image=ALPINE_IMAGE, docker_poll_interval=1.0, ) - logging.info("Scout module started") - s = scout.Scout(cfg, callback) - logging.info("Submitting task to Scout module") - - if args.test1: - seconds = args.test1 - task_1(s, seconds) - if args.test2: - seconds = args.test2 - task_2(s, seconds) + logging.info("Running tests") + for test in args.tests: + TEST_FUNCTIONS[ test - 1 ](cfg) - if args.test3: - seconds = args.test3 - task_3(s, seconds) + if args.run_scout: + run_scout(cfg) - if args.scout_suite: - test_scout(s) + logging.info("Done") if __name__ == "__main__":