Created 3 tests for containers

scout/scout.py

@@ -15,14 +15,15 @@
 
 import docker
 
+from typing import List
 
 DOCKER_STATUSCODE_KEY = "StatusCode"
 OUTDIR_CONTAINER_MOUNT = "/root/output"
 SCOUT_TASK_LABEL_KEY = "scout-task-id"
 
 
-# https://docs.docker.com/engine/reference/commandline/ps/
-class ContainerState(enum.StrEnum):
+# # https://docs.docker.com/engine/reference/commandline/ps/
+class ContainerState(enum.Enum):
     CREATED = "created"
     RESTARTING = "restarting"
     RUNNING = "running"
@@ -34,17 +35,30 @@ class ContainerState(enum.StrEnum):
     def is_done(self):
         return self in [ContainerState.EXITED, ContainerState.DEAD]
 
+# Set de configuração do scout especificamente
+# @dataclasses.dataclass(frozen=True)
+# class ScoutConfig:
+#     credentials_file: pathlib.Path
+#     output_dir: pathlib.Path
+#     docker_image: str = "rossja/ncc-scoutsuite:aws-latest"
+#     docker_poll_interval: float = 16.0
+#     docker_socket: str | None = None
+#     docker_timeout: int = 5
+    # definir os comandos da imagem para o alpine
+
+# commands scout - usado para rodar o modulo
+# commands_scout = ["scout","aws","--no-browser","--result-format","json","--report-dir",f"{OUTDIR_CONTAINER_MOUNT}","--logfile",f"{OUTDIR_CONTAINER_MOUNT}/scout.log",]
+# volumes_scout = {str(self.config.credentials_file): {"bind": "/root/.aws/credentials","mode": "ro",},str(outfp): {"bind": OUTDIR_CONTAINER_MOUNT,"mode": "rw",},}
 
 @dataclasses.dataclass(frozen=True)
 class ScoutConfig:
     credentials_file: pathlib.Path
     output_dir: pathlib.Path
-    docker_image: str = "rossja/ncc-scoutsuite:aws-latest"
+    docker_image: str = "alpine"
     docker_poll_interval: float = 16.0
     docker_socket: str | None = None
     docker_timeout: int = 5
 
-
 class Task(Protocol):
     label: str
 
@@ -68,7 +82,6 @@ def enqueue(self, taskcfg: Task) -> None:
     def shutdown(self, wait: bool) -> None:
         ...
 
-
 class Scout(ScanModule):
     def __init__(self, config: ScoutConfig, callback: TaskCompletionCallback) -> None:
         def get_docker_client() -> docker.DockerClient:
@@ -82,43 +95,29 @@ def get_docker_client() -> docker.DockerClient:
         self.containers: set = set()
         self.task_completion_callback = callback
         self.lock: threading.Lock = threading.Lock()
-        self.thread: threading.Thread = threading.Thread(
-            target=self.scout_polling_thread,
-            name="scout-polling-thread",
-        )
+        self.thread: threading.Thread = threading.Thread(target=self.scout_polling_thread,name="scout-polling-thread",)
         self.thread.start()
 
-    def enqueue(self, taskcfg: Task) -> None:
+    def enqueue(self, taskcfg: Task, commands: List[str], time: int = 0) -> None:
         assert isinstance(taskcfg, ScoutTask)
+
+        if(time != 0):
+            commands[1] = f"{time}"
+
         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=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",
-                    },
+                    #####
                 },
                 working_dir="/root",
             )
@@ -132,14 +131,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 +153,38 @@ 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")
+                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",
+                logging.info("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,
+                self.config.docker_poll_interval
             )
+
         for ctx, _cfg in completed:
-            ctx.remove()
+            ctx.remove()

scout/tests/testrun.py

@@ -8,6 +8,14 @@
 
 import scout
 
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-t1', '--task1', type=int ,help='Task 1')
+parser.add_argument('-t2', '--task2', type=int , help='Task 2')
+parser.add_argument('-t3', '--task3', type=int , help='Task 3')
+args = parser.parse_args()
+
 
 CWD = pathlib.Path(os.getcwd())
 assert pathlib.Path(CWD / __file__).exists(), "Run from inside tests/ with ./testrun.py"
@@ -17,14 +25,39 @@
 def callback(label: str, success: bool) -> None:
     logging.info("Callback for %s running, success=%s", label, success)
 
+def task_1(s, taskcfg, seconds: int):
+    task_generic_sleep = ["sleep", f'{1}']
+    logging.info("")
+    logging.info("RUNNING TASK 1")
+    s.enqueue(taskcfg, task_generic_sleep, seconds)
+    logging.info("Task submitted")
+    s.shutdown()
+
+def task_2(s, taskcfg, seconds: int):
+    task_generic_inf = ["tail", "-f" ,"/dev/null"]
+    logging.info("")
+    logging.info("RUNNING TASK 2")
+    s.enqueue(taskcfg, task_generic_inf)
+    time.sleep(seconds)
+    logging.info("Task submitted")
+    s.shutdown(False)
+
+def task_3(s, taskcfg, seconds: int):
+    task_generic_error = ["sh", "-c", f'sleep {seconds}', "&&" , "exit", "1"]
+    logging.info("")
+    logging.info("RUNNING TASK 3")
+    s.enqueue(taskcfg, task_generic_error)
+    logging.info("Task submitted")
+    s.shutdown()
+
+
 
 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_poll_interval=1.0,
     )
     logging.info("Scout module started")
@@ -33,9 +66,21 @@ def main():
     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.task1:
+        seconds = args.task1
+        task_1(s, taskcfg, seconds)
+
+    #python3 testrun.py -t2 "tempo de execução em segundos"
+    if args.task2:
+        seconds = args.task2
+        task_2(s, taskcfg, seconds)
+
+    #python3 testrun.py -t3 "tempo de execução em segundos"
+    if args.task3:
+        seconds = args.task3
+        task_3(s, taskcfg, seconds)
 
 
 if __name__ == "__main__":