Merge pull request #14 from SpecterOps/monitor_mode

add 'monitor' command to submit_to_nemesis

cmd/dashboard/dashboard/Nemesis.py

@@ -39,11 +39,11 @@ def build_page(username: str):
     st.subheader("Files")
     cols = st.columns(5)
     with cols[0]:
-        num_plaintext_documents = utils.get_elastic_total_indexed_documents("file_data_plaintext")
-        st.metric("Indexed Documents", num_plaintext_documents)
-    with cols[1]:
         num_enriched_documents = utils.get_elastic_total_indexed_documents("file_data_enriched")
         st.metric("Processed Files", num_enriched_documents)
+    with cols[1]:
+        num_plaintext_documents = utils.get_elastic_total_indexed_documents("file_data_plaintext")
+        st.metric("Indexed Documents", num_plaintext_documents)
     with cols[2]:
         num_np_matches = utils.get_elastic_total_indexed_documents("file_data_enriched", query={"exists": {"field": "noseyparker"}})
         st.metric("NoseyParker Matches", num_np_matches)

cmd/enrichment/.vscode/launch.json

@@ -55,6 +55,26 @@
                 "--folder",
                 "${workspaceFolder}/../../sample_files"
             ]
+        },
+        {
+            "name": "submit_to_nemesis.py - Monitor",
+            "type": "python",
+            "request": "launch",
+            "module": "enrichment.cli.submit_to_nemesis",
+            "console": "integratedTerminal",
+            "cwd": "${workspaceFolder}",
+            "justMyCode": false,
+            "env": {
+                // "PYTHONASYNCIODEBUG": "1",
+                "BETTER_EXCEPTIONS": "1",
+                "FORCE_COLOR": "1",
+            },
+            "args": [
+                "--monitor",
+                "/tmp/mon",
+                // "-l",
+                // "DEBUG"
+            ]
         }
     ]
 }

cmd/enrichment/enrichment/cli/submit_to_nemesis/__main__.py

@@ -44,7 +44,7 @@ def main():
     loop.set_exception_handler(handle_exception)
 
     try:
-        task = loop.create_task(amain())
+        task = loop.create_task(amain(loop))
         loop.run_until_complete(task)
     finally:
         loop.close()

cmd/enrichment/enrichment/cli/submit_to_nemesis/file_monitoring.py

@@ -0,0 +1,39 @@
+# Standard Libraries
+import asyncio
+import os
+import time
+
+# 3rd Party Libraries
+import structlog
+from watchdog.events import FileSystemEventHandler
+from watchdog.observers import Observer
+
+logger = structlog.getLogger()
+
+
+class NewFileHandler(FileSystemEventHandler):
+    def __init__(self, queue, loop: asyncio.AbstractEventLoop):
+        self.queue = queue
+        self.loop = loop
+
+    def on_created(self, event):
+        if event.is_directory:
+            return
+
+        self.loop.call_soon_threadsafe(self.loop.create_task, self.queue.put(event.src_path))
+
+
+async def monitor_directory(directory, loop):
+    queue = asyncio.Queue()
+    event_handler = NewFileHandler(queue, loop)
+    observer = Observer()
+    observer.schedule(event_handler, directory, recursive=True)
+    observer.start()
+
+    try:
+        while True:
+            file_path = await queue.get()
+            yield file_path
+    finally:
+        observer.stop()
+        observer.join()

cmd/enrichment/enrichment/cli/submit_to_nemesis/submit_to_nemesis.py

@@ -20,6 +20,7 @@ import requests
 import structlog
 import urllib3
 import yaml
+from enrichment.cli.submit_to_nemesis.file_monitoring import monitor_directory
 from nemesiscommon.apiclient import FileUploadRequest, NemesisApiClient
 from nemesiscommon.logging import configure_logger
 from structlog.typing import FilteringBoundLogger
@@ -39,6 +40,7 @@ async def get_config() -> dict[str, str]:
     parser = argparse.ArgumentParser(description="Submit file(s) to Nemesis.", prog="submit_to_nemesis")
     parser.add_argument("-f", "--file", type=str, nargs="+", help="File(s) to submit to Nemesis")
     parser.add_argument("--folder", type=str, nargs="+", help="Folders(s) to submit to Nemesis")
+    parser.add_argument("-m", "--monitor", type=str, nargs="?", help="Folder to monitor for new files")
     parser.add_argument("-s", "--sec_between_files", type=float, nargs="?", default=0, help="Seconds between file submissions (default 0). If > 0, it cannot be used with the --workers argument.")
     parser.add_argument("-w", "--workers", type=int, nargs="?", default=10, help="Number of workers to use (default 10). If --sec_between_files argument is set, value will be 1")
     parser.add_argument("-r", "--repeat", type=int, default=0, help="Times to repeat the submission (for stress testing)")
@@ -115,6 +117,7 @@ async def get_config() -> dict[str, str]:
 
     config["file"] = args.file
     config["folder"] = args.folder
+    config["monitor"] = args.monitor
     config["sec_between_files"] = args.sec_between_files
     config["repeat"] = args.repeat
     config["timeout"] = args.timeout
@@ -263,18 +266,21 @@ async def submit_random_cookies(config, num_cookies=1000) -> uuid.UUID | None:
     for i in range(num_cookies):
         if i % 100 == 0:
             domain = f"{'-'.join(random.choices(words, k=2))}.com"
-        cookie_data.append({"user_data_directory": "C:/Users/harmj0y/AppData/Local/Google/Chrome/User Data/Default/Cookies",
-                            "domain": domain,
-                            "path": "/",
-                            "name": f"VALUE_{i}",
-                            "value": f"{random.choice(words)}_{random.randint(1, 10000000)}",
-                            "expires": "2030-01-01T01:01:01.000Z",
-                            "secure": True,
-                            "http_only": True,
-                            "session": False,
-                            "samesite": "lax",
-                            "source_port": 443
-                            })
+        cookie_data.append(
+            {
+                "user_data_directory": "C:/Users/harmj0y/AppData/Local/Google/Chrome/User Data/Default/Cookies",
+                "domain": domain,
+                "path": "/",
+                "name": f"VALUE_{i}",
+                "value": f"{random.choice(words)}_{random.randint(1, 10000000)}",
+                "expires": "2030-01-01T01:01:01.000Z",
+                "secure": True,
+                "http_only": True,
+                "session": False,
+                "samesite": "lax",
+                "source_port": 443,
+            }
+        )
 
     resp = await nemesis_post_data(config, {"metadata": metadata, "data": cookie_data})
     return uuid.UUID(resp["object_id"]) if resp else None
@@ -286,7 +292,7 @@ async def process_file(config, file_path) -> uuid.UUID | None:
     then posting the file_data message.
     """
 
-    if "services_api.json" in file_path:  # NOTE: This must not conflict with the example Seatbelt services either
+    if re.match(r".*services_api.*\.json$", file_path):  # NOTE: This must not conflict with the example Seatbelt services either
         with open(file_path, "r") as f:
             services_json_raw = f.read()
             services_json = json.loads(services_json_raw)
@@ -536,6 +542,10 @@ def return_args_and_exceptions(func, exception_handler: Callable) -> Callable:
     return functools.partial(_return_args_and_exceptions, func)
 
 
+def exception_handler(e, args):
+    logger.exception("Error processing file", args=args)
+
+
 async def submit_paths_concurrently(config, paths: List[str], workers: int, delay: float = 0) -> AsyncIterator[Tuple[str, uuid.UUID]]:
     """Submits files to Nemesis concurrently.
 
@@ -598,9 +608,6 @@ async def submit_paths_concurrently(config, paths: List[str], workers: int, dela
 
         return file_uuid
 
-    def exception_handler(e, args):
-        logger.exception("Error processing file", args=args)
-
     wrapped_process_file = return_args_and_exceptions(process_file_local, exception_handler)
 
     try:
@@ -613,6 +620,82 @@ async def submit_paths_concurrently(config, paths: List[str], workers: int, dela
     logger.info(f"Completed processing {result_count} files out of {total_file_count} total files")
 
 
+async def wait_for_stable_size(file_path, delay=1.0, retries=600):
+    logger.debug("Waiting for file size to stabilize", path=file_path)
+
+    previous_size = -1
+    current_size = 0
+    tries = 0
+    while tries < retries:
+        try:
+            current_size = os.path.getsize(file_path)
+            if current_size == previous_size:
+                return True
+            previous_size = current_size
+            await asyncio.sleep(delay)
+            tries += 1
+        except Exception as e:
+            logger.error("Error accessing a newly create file", path=file_path, exception=e)
+            return False
+    return False
+
+
+async def monitor_submit_paths_concurrently(config, path_iter: AsyncIterator, workers: int, delay: float = 0) -> AsyncIterator[Tuple[str, uuid.UUID]]:
+    """Submits files to Nemesis concurrently.
+
+    Args:
+        config (_type_): submit_to_nemesis configuration object.
+        paths (List[str]): List of file or folder paths to submit to Nemesis.
+        workers (int): Number of concurrent tasks that can run at once.
+        delay (float, optional): Time delay between each submission. Defaults to 0.
+
+    Returns:
+        AsyncIterator[Tuple[str, uuid.UUID]]: _description_
+
+    Yields:
+        Iterator[AsyncIterator[Tuple[str, uuid.UUID]]]: _description_
+    """
+    result_count = 0
+
+    async def monitor_process_file_local(file_path) -> uuid.UUID | None:
+        global logger
+
+        structlog.contextvars.bind_contextvars(
+            file_path=file_path,
+        )
+
+        logger.info("Processing file")
+
+        if not await wait_for_stable_size(file_path):
+            logger.error("An error occurred while waiting for the file. File did not process.", path=file_path)
+            return
+
+        file_uuid = await process_file(config, file_path)
+
+        logger.info(
+            f"Done processing file {result_count+1}",
+            file_uuid=str(file_uuid),
+            total_completed=result_count + 1,
+        )
+        structlog.contextvars.clear_contextvars()
+
+        if delay > 0:
+            await asyncio.sleep(delay)
+
+        return file_uuid
+
+    wrapped_process_file = return_args_and_exceptions(monitor_process_file_local, exception_handler)
+
+    try:
+        async for result in map_unordered(wrapped_process_file, path_iter, limit=workers):
+            result_count += 1
+            yield result
+    except asyncio.CancelledError:
+        logger.warn("Cancelled file uploads")
+
+    logger.info(f"Completed processing {result_count} total files")
+
+
 async def is_file_processed(es_client, message_id: int):
     query = {"bool": {"filter": [{"match_phrase": {"metadata.messageId": message_id}}]}}
 
@@ -628,9 +711,47 @@ async def is_file_processed(es_client, message_id: int):
         return False
 
 
-async def amain():
+async def submit_files_and_folders(config: dict[str, str]):
     paths_to_process = []
     processed_file_uuids = []
+
+    logger.info("Submitting files/folders to Nemesis")
+    if config["file"]:
+        for f in config["file"]:
+            paths_to_process.append(f)
+
+    if config["folder"]:
+        for f in config["folder"]:
+            paths_to_process.append(f)
+
+    for i in range(config["repeat"] + 1):
+        logger.info("Waiting for tasks to complete")
+        async for result in submit_paths_concurrently(config, paths_to_process, config["workers"], config["sec_between_files"]):
+            path, file_uuid = result
+            processed_file_uuids.append(file_uuid)
+
+
+async def monitor_and_submit_folder_files(config: dict[str, str], loop):
+    processed_file_uuids = []
+
+    path = os.path.abspath(config["monitor"])
+
+    if not os.path.exists(path):
+        logger.error("Path does not exist", path=path)
+        return
+
+    if os.path.isfile(path):
+        logger.error("The monitor path is a file, not a folder", path=path)
+
+    logger.info("Monitoring a folder for new files to submit", path=path)
+
+    iter = monitor_directory(path, loop)
+    async for result in monitor_submit_paths_concurrently(config, iter, config["workers"], config["sec_between_files"]):
+        path, file_uuid = result
+        processed_file_uuids.append(file_uuid)
+
+
+async def amain(loop):
     config = await get_config()
     if not config:
         return
@@ -638,20 +759,11 @@ async def amain():
     try:
         if config["cookies"]:
             await submit_random_cookies(config, config["cookies"])
-        else:
-            if config["file"]:
-                for f in config["file"]:
-                    paths_to_process.append(f)
+        elif config["file"] or config["folder"]:
+            await submit_files_and_folders(config)
+        elif config["monitor"]:
+            await monitor_and_submit_folder_files(config, loop)
 
-            if config["folder"]:
-                for f in config["folder"]:
-                    paths_to_process.append(f)
-
-            for i in range(config["repeat"] + 1):
-                logger.info("Waiting for tasks to complete")
-                async for result in submit_paths_concurrently(config, paths_to_process, config["workers"], config["sec_between_files"]):
-                    path, file_uuid = result
-                    processed_file_uuids.append(file_uuid)
     except asyncio.CancelledError:
         pass
 

cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py

@@ -92,10 +92,10 @@ async def get_all_rabbit_mq_queues(rabbit_mq_api_url: str) -> List[str]:
             if r.status_code == 200:
                 return [queue["name"] for queue in r.json()]
             else:
-                logger.aerror("Error retrieving RabbitMQ queues from the API", status=r.status_code)
+                await logger.aerror("Error retrieving RabbitMQ queues from the API", status=r.status_code)
                 return []
         except Exception as e:
-            logger.aerror("Error retrieving RabbitMQ queues from the API", exception=e)
+            await logger.aerror("Error retrieving RabbitMQ queues from the API", exception=e)
             return []
 
 
@@ -453,5 +453,5 @@ class NemesisApiRoutes(Routable):
                 #   ref - https://github.com/tiangolo/fastapi/issues/2152#issuecomment-889282903
                 return FileResponse(file.name, background=BackgroundTask(os.remove, file.name), media_type=content_type, headers=headers)
         except Exception as e:
-            logger.aexception(e, message="Failed to download file", file_uuid=id)
+            await logger.aerror(message="Failed to download file", file_uuid=id, exception=e)
             return Response(status_code=404, content="File not found")

kubernetes/pgadmin/statefulset.yaml

@@ -36,6 +36,8 @@ spec:
                   key: pgadmin-password
             - name: SCRIPT_NAME
               value: /pgadmin/
+            - name: MAX_LOGIN_ATTEMPTS
+              value: "10"
           ports:
             - name: http
               containerPort: 80

sample_files/structured/services_api.json

@@ -27,7 +27,7 @@
         "agent_id": "339429212",
         "agent_type": "beacon",
         "automated": false,
-        "data_type": "services",
+        "data_type": "service",
         "expiration": "2024-04-03T10:08:40.000Z",
         "source": "DC",
         "project": "ASSESS-X",