diff --git a/cmd/dashboard/dashboard/Nemesis.py b/cmd/dashboard/dashboard/Nemesis.py index 1f057db..605c5f9 100644 --- a/cmd/dashboard/dashboard/Nemesis.py +++ b/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) diff --git a/cmd/enrichment/.vscode/launch.json b/cmd/enrichment/.vscode/launch.json index 23be6d2..76578e4 100644 --- a/cmd/enrichment/.vscode/launch.json +++ b/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" + ] } ] } \ No newline at end of file diff --git a/cmd/enrichment/enrichment/cli/submit_to_nemesis/__main__.py b/cmd/enrichment/enrichment/cli/submit_to_nemesis/__main__.py index 85a86cf..feed3ce 100644 --- a/cmd/enrichment/enrichment/cli/submit_to_nemesis/__main__.py +++ b/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() diff --git a/cmd/enrichment/enrichment/cli/submit_to_nemesis/file_monitoring.py b/cmd/enrichment/enrichment/cli/submit_to_nemesis/file_monitoring.py new file mode 100644 index 0000000..6da0748 --- /dev/null +++ b/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() diff --git a/cmd/enrichment/enrichment/cli/submit_to_nemesis/submit_to_nemesis.py b/cmd/enrichment/enrichment/cli/submit_to_nemesis/submit_to_nemesis.py index f141701..3389618 100644 --- a/cmd/enrichment/enrichment/cli/submit_to_nemesis/submit_to_nemesis.py +++ b/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 diff --git a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py b/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py index 1add358..f2d56fb 100644 --- a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py +++ b/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") diff --git a/kubernetes/pgadmin/statefulset.yaml b/kubernetes/pgadmin/statefulset.yaml index 88cb499..9a21c35 100644 --- a/kubernetes/pgadmin/statefulset.yaml +++ b/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 diff --git a/sample_files/structured/services.json b/sample_files/structured/services_api.json similarity index 98% rename from sample_files/structured/services.json rename to sample_files/structured/services_api.json index ce1cfe7..5411290 100644 --- a/sample_files/structured/services.json +++ b/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",