Created 3 tests for containers

scout/scout.py

@@ -22,7 +22,7 @@
 
 
 # https://docs.docker.com/engine/reference/commandline/ps/
-class ContainerState(enum.StrEnum):
+class ContainerState(enum.Enum):
     CREATED = "created"
     RESTARTING = "restarting"
     RUNNING = "running"
@@ -39,7 +39,7 @@ def is_done(self):
 class ScoutConfig:
     credentials_file: pathlib.Path
     output_dir: pathlib.Path
-    docker_image: str = "rossja/ncc-scoutsuite:aws-latest"
+    docker_image: str
     docker_poll_interval: float = 16.0
     docker_socket: str | None = None
     docker_timeout: int = 5
@@ -52,6 +52,8 @@ class Task(Protocol):
 @dataclasses.dataclass(frozen=True)
 class ScoutTask:
     label: str
+    commands: str | None = None
+    volumes: dict | None = None
     aws_api_key: str | None = None
     aws_api_secret: str | None = None
     role_arn: str | None = None
@@ -63,8 +65,10 @@ class ScoutTask:
 class ScanModule(Protocol):
     def __init__(self, config: Any, callback: TaskCompletionCallback) -> None:
         ...
+
     def enqueue(self, taskcfg: Task) -> None:
         ...
+
     def shutdown(self, wait: bool) -> None:
         ...
 
@@ -90,36 +94,18 @@ 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:
             ctx = self.docker.containers.run(
                 self.config.docker_image,
-                command=[
-                    "scout",
-                    "aws",
-                    "--no-browser",
-                    "--result-format",
-                    "json",
-                    "--report-dir",
-                    f"{OUTDIR_CONTAINER_MOUNT}",
-                    "--logfile",
-                    f"{OUTDIR_CONTAINER_MOUNT}/scout.log",
-                ],
+                command=taskcfg.commands,
                 detach=True,
                 labels={SCOUT_TASK_LABEL_KEY: taskcfg.label},
                 stdout=True,
                 stderr=True,
-                volumes={
-                    str(self.config.credentials_file): {
-                        "bind": "/root/.aws/credentials",
-                        "mode": "ro",
-                    },
-                    str(outfp): {
-                        "bind": OUTDIR_CONTAINER_MOUNT,
-                        "mode": "rw",
-                    },
-                },
+                volumes=taskcfg.volumes,
                 working_dir="/root",
             )
         except docker.errors.APIError as e:
@@ -132,14 +118,16 @@ 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)
-        logging.info("Joining Scout polling thread")
-        self.thread.join()
+            self.containers.clear()
+        else:
+            self.handle_finished_containers()
+            self.thread.join()
+
         logging.info("Joined Scout polling thread, module shut down")
 
     def scout_polling_thread(self) -> None:
@@ -152,32 +140,43 @@ def handle_finished_containers(self) -> None:
         completed = set()
         with self.lock:
             for ctx, cfg in self.containers:
-                ctx.reload()
+                try:
+                    ctx.reload()
+                except docker.errors.NotFound:
+                    logging.warning("Container not found: %s", cfg.label)
+                    continue
+
                 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
                 ).decode("utf8")
                 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)
                 self.task_completion_callback(cfg.label, True)
                 logging.info(
-                    "Scout run completed, id %s status %d",
+                    "Scout run completed, id %s sta+-tus %d",
                     cfg.label,
                     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()

scout/tests/testrun.py

@@ -8,34 +8,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()
+
 
 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 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",  # Commands_scout
+        volumes={
+            "./dev/aws-credentials": {"bind": "/root/.aws/credentials", "mode": "ro"},
+            "./data": {"bind": "/root/output", "mode": "rw"},
+        },  # volumes_scout
+    )
+
+    logging.info("RUNNING TEST SCOUT")
+    s.enqueue(taskcfg)
+    logging.info("Task submitted")
+    s.shutdown()
+
+
+def task_1(s, seconds: int):
+    taskcfg = scout.ScoutTask(
+        time.strftime("scout-%Y%m%d-%H%M%S"),
+        f"sleep {seconds}",  # commands
+    )
+
+    taskcfg2 = scout.ScoutTask(time.strftime("scout-%Y%m%d-%H%M%S"), f"sleep 5")
+
+    logging.info(
+        "RUNNING TEST 1 - The container was removed after n seconds, along with another set of tasks."
+    )
+    s.enqueue(taskcfg)
+    s.enqueue(taskcfg2)
+    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",  # commands
+    )
+
+    logging.info(
+        "RUNNING TEST 2 - Container with infinite execution removed after n seconds."
+    )
+    s.enqueue(taskcfg)
+
+    time.sleep(seconds)
+    logging.info("Task submitted")
+    s.shutdown(False)
+
+
+def task_3(s, seconds: int):
+    taskcfg = scout.ScoutTask(
+        time.strftime("scout-%Y%m%d-%H%M%S"),
+        f"sh -c 'sleep {seconds} && exit 1'",
+    )
+
+    logging.info("RUNNING TEST 3 - Container removed after returning an error.")
+    s.enqueue(taskcfg)
+
+    logging.info("Task submitted")
+    s.shutdown()
+
+
+scout_image = "rossja/ncc-scoutsuite:aws-latest"
+generic_image = "alpine"
+
+
 def main():
     logging.basicConfig(level=logging.INFO)
     logging.info("Starting Scout module")
     cfg = scout.ScoutConfig(
         pathlib.Path(SCOUT_MOD_PATH / "dev/aws-credentials"),
         pathlib.Path(SCOUT_MOD_PATH / "data"),
-        # docker_image="hello-world",
+        docker_image=scout_image,
         docker_poll_interval=1.0,
     )
     logging.info("Scout module started")
     s = scout.Scout(cfg, callback)
     logging.info("Submitting task to Scout module")
-    taskcfg = scout.ScoutTask(
-        time.strftime("scout-%Y%m%d-%H%M%S"),
-    )
-    s.enqueue(taskcfg)
-    logging.info("Task submitted")
-    s.shutdown()
+
+    # python3 testrun.py -t1 "tempo de execução em segundos"
+    if args.test1:
+        seconds = args.test1
+        task_1(s, seconds)
+
+    # #python3 testrun.py -t2 "tempo de execução em segundos"
+    if args.test2:
+        seconds = args.test2
+        task_2(s, seconds)
+
+    # python3 testrun.py -t3 "tempo de execução em segundos"
+    if args.test3:
+        seconds = args.test3
+        task_3(s, seconds)
+
+    # python3 testrun.py -scout
+    if args.scout_suite:
+        test_scout(s)
 
 
 if __name__ == "__main__":