From 14ae40991200e7078c0b1b9bf94421c47a8e99f5 Mon Sep 17 00:00:00 2001 From: Will Date: Fri, 8 Mar 2024 00:21:12 -0800 Subject: [PATCH 1/2] /reset route now clears minio files as well /reset route now clears minio files as well --- .../enrichment/tasks/webapi/nemesis_api.py | 13 +++++++------ .../python/nemesiscommon/nemesiscommon/storage.py | 4 ++++ .../nemesiscommon/nemesiscommon/storage_minio.py | 12 ++++++++++++ .../nemesiscommon/nemesiscommon/storage_s3.py | 5 +++++ 4 files changed, 28 insertions(+), 6 deletions(-) diff --git a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py b/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py index 4ae99cf..83baa37 100644 --- a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py +++ b/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py @@ -17,15 +17,12 @@ import uvicorn from aio_pika import connect_robust from elasticsearch import AsyncElasticsearch from enrichment.cli.submit_to_nemesis.submit_to_nemesis import ( - map_unordered, - return_args_and_exceptions, -) + map_unordered, return_args_and_exceptions) from enrichment.lib.nemesis_db import NemesisDb from enrichment.lib.registry import include_registry_value -from fastapi import FastAPI, Request, APIRouter, HTTPException +from fastapi import APIRouter, FastAPI, HTTPException, Request from fastapi.responses import FileResponse, Response from google.protobuf.json_format import Parse - # from nemesiscommon.clearqueues import clearRabbitMQQueues from nemesiscommon.constants import ALL_ES_INDICIES, NemesisQueue from nemesiscommon.messaging import MessageQueueProducerInterface @@ -243,7 +240,7 @@ class NemesisApiRoutes(): return Response() async def reset(self): - """When called, purges Postgres, Elastic, and RabbitMQ.""" + """When called, purges Postgres, Elastic, RabbitMQ, and datalake files.""" await logger.ainfo("Clearing datastore!") # first purge all RabbitMQ queue messages @@ -261,6 +258,10 @@ class NemesisApiRoutes(): await logger.ainfo("Clearing Elastic index", index=ES_INDEX) await self.es_client.indices.delete(index=ES_INDEX) + # and finally clear the files from the datalake + await logger.ainfo("Deleting files in the datalake.") + await self.storage.delete_all_files() + async def reprocess(self): """When called, triggers the reprocessing of all existing data messages.""" diff --git a/packages/python/nemesiscommon/nemesiscommon/storage.py b/packages/python/nemesiscommon/nemesiscommon/storage.py index 493b3b6..be7ac27 100644 --- a/packages/python/nemesiscommon/nemesiscommon/storage.py +++ b/packages/python/nemesiscommon/nemesiscommon/storage.py @@ -19,6 +19,10 @@ class StorageInterface(ABC): async def exists(self, file_name: str) -> bool: raise NotImplementedError + @abstractmethod + async def delete_all_files(self) -> bool: + raise NotImplementedError + @abstractmethod async def __aenter__(self): raise NotImplementedError diff --git a/packages/python/nemesiscommon/nemesiscommon/storage_minio.py b/packages/python/nemesiscommon/nemesiscommon/storage_minio.py index 17c2aa0..46a4b86 100644 --- a/packages/python/nemesiscommon/nemesiscommon/storage_minio.py +++ b/packages/python/nemesiscommon/nemesiscommon/storage_minio.py @@ -55,6 +55,18 @@ class StorageMinio(StorageInterface): async def exists(self, file_name: str) -> bool: raise NotImplementedError + async def delete_all_files(self) -> bool: + await logger.adebug("Deleting all files from bucket", bucket_name=self.assessment_id) + + try: + files = await self.minio_client.list_objects(self.assessment_id, recursive=True) + for file in files: + await self.minio_client.remove_object(self.assessment_id, file.object_name) + return True + except Exception as e: + await logger.aexception(e, message="Failed to delete files from bucket", bucket_name=self.assessment_id) + raise + async def __aenter__(self): return self diff --git a/packages/python/nemesiscommon/nemesiscommon/storage_s3.py b/packages/python/nemesiscommon/nemesiscommon/storage_s3.py index e56722a..d591dc0 100644 --- a/packages/python/nemesiscommon/nemesiscommon/storage_s3.py +++ b/packages/python/nemesiscommon/nemesiscommon/storage_s3.py @@ -90,6 +90,11 @@ class StorageS3(StorageInterface): async def exists(self, file_name: str) -> bool: raise NotImplementedError + async def delete_all_files(self) -> bool: + await logger.awarning("Deleting all files from S3 not yet implemented!") + # async with aioboto3.Session().client(service_name="s3", **self._aws_client_args) as s3_client: + # s3_client. + async def get_s3_file_path(self, file_uuid: uuid.UUID) -> str: return f"{self.assessment_id}/{file_uuid}.enc" From 76559874b1fa3595da50bdfb01553e1000479a11 Mon Sep 17 00:00:00 2001 From: Will Date: Fri, 8 Mar 2024 11:07:57 -0800 Subject: [PATCH 2/2] Added Minio storage expiration policy for bucket on creation - Added Minio storage expiration policy for bucket on creation --- cmd/enrichment/enrichment/containers.py | 1 + cmd/enrichment/enrichment/settings.py | 1 + .../enrichment/tasks/webapi/nemesis_api.py | 10 ++++++- .../nemesiscommon/storage_minio.py | 29 ++++++++++++++++++- 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/cmd/enrichment/enrichment/containers.py b/cmd/enrichment/enrichment/containers.py index 662dd9b..80864af 100644 --- a/cmd/enrichment/enrichment/containers.py +++ b/cmd/enrichment/enrichment/containers.py @@ -594,6 +594,7 @@ class Container(containers.DeclarativeContainer): config.assessment_id, config.log_level, config.reprocessing_workers, + config.storage_expiration_days, ) task_landingpage = providers.Factory(LandingPageApi, config.log_level) task_yara_api = providers.Factory(YaraApi, storage_service, config.yara_api_port, config.data_download_dir, config.log_level) diff --git a/cmd/enrichment/enrichment/settings.py b/cmd/enrichment/enrichment/settings.py index c6cee98..647a0df 100644 --- a/cmd/enrichment/enrichment/settings.py +++ b/cmd/enrichment/enrichment/settings.py @@ -66,6 +66,7 @@ class EnrichmentSettings(FileProcessingService): # type: ignore reprocessing_workers: PositiveInt = 5 # number of parallel reprocessing workers tasks: Optional[List[str]] # List of tasks to start registry_value_batch_size: PositiveInt = 5000 # number of registry values to emit per reg carving message + storage_expiration_days: PositiveInt = 100 # number of days to set for file storage auto-expiry @validator("slack_channel") def slack_channel_is_valid(cls, channel: Optional[str]) -> Optional[str]: diff --git a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py b/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py index 83baa37..c3ab921 100644 --- a/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py +++ b/cmd/enrichment/enrichment/tasks/webapi/nemesis_api.py @@ -122,6 +122,7 @@ class NemesisApi(TaskInterface): assessment_id: str log_level: str reprocessing_workers: int + storage_expiration_days: int def __init__( self, @@ -134,6 +135,7 @@ class NemesisApi(TaskInterface): assessment_id: str, log_level: str, reprocessing_workers: int, + storage_expiration_days: int ) -> None: self.storage = storage self.rabbitmq_connection_uri = rabbitmq_connection_uri @@ -144,6 +146,7 @@ class NemesisApi(TaskInterface): self.assessment_id = assessment_id self.log_level = log_level self.reprocessing_workers = reprocessing_workers + self.storage_expiration_days = storage_expiration_days async def run(self) -> None: app = FastAPI(title="Nemesis API") @@ -156,6 +159,7 @@ class NemesisApi(TaskInterface): self.queue_map, self.assessment_id, self.reprocessing_workers, + self.storage_expiration_days, ) app.include_router(routes.router) @@ -199,6 +203,7 @@ class NemesisApiRoutes(): producers: Dict[NemesisQueue, MessageQueueProducerInterface] assessment_id: str reprocessing_workers: int + storage_expiration_days: int def __init__( self, @@ -210,6 +215,7 @@ class NemesisApiRoutes(): queues: Dict[NemesisQueue, MessageQueueProducerInterface], assessment_id: str, reprocessing_workers: int, + storage_expiration_days: int, ) -> None: super().__init__() self.storage = storage @@ -220,6 +226,7 @@ class NemesisApiRoutes(): self.producers = queues self.assessment_id = assessment_id self.reprocessing_workers = reprocessing_workers + self.storage_expiration_days = storage_expiration_days self.router = APIRouter() self.router.add_api_route("/", self.home, methods=["GET"]) self.router.add_api_route("/ready", self.ready, methods=["GET"]) @@ -395,7 +402,8 @@ class NemesisApiRoutes(): with open(tmpfile.name, "wb") as f: f.write(await request.body()) - file_uuid = await self.storage.upload(f.name) + # the first file that comes in will set the bucket file expiry policy (for now) + file_uuid = await self.storage.upload(f.name, self.storage_expiration_days) return {"object_id": str(file_uuid)} @aio.time(Summary("download", "Download file")) # type: ignore diff --git a/packages/python/nemesiscommon/nemesiscommon/storage_minio.py b/packages/python/nemesiscommon/nemesiscommon/storage_minio.py index 46a4b86..6795daa 100644 --- a/packages/python/nemesiscommon/nemesiscommon/storage_minio.py +++ b/packages/python/nemesiscommon/nemesiscommon/storage_minio.py @@ -7,6 +7,8 @@ from typing import Optional, Type # 3rd Party Libraries import structlog from miniopy_async import Minio +from miniopy_async.commonconfig import ENABLED, Filter +from miniopy_async.lifecycleconfig import LifecycleConfig, Rule, Expiration from nemesiscommon.storage import StorageInterface logger = structlog.get_logger(module=__name__) @@ -44,9 +46,26 @@ class StorageMinio(StorageInterface): return temp_file - async def upload(self, file_path: str) -> uuid.UUID: + async def upload(self, file_path: str, storage_expiration_days: int = 100) -> uuid.UUID: if not await self.minio_client.bucket_exists(self.assessment_id): + await logger.ainfo("Creating Minio bucket", bucket=self.assessment_id) await self.minio_client.make_bucket(self.assessment_id) + + # since this is the only place that creates the bucket, we can set + # the auto-expiration policy here + config = LifecycleConfig( + [ + Rule( + ENABLED, + rule_filter=Filter(prefix=""), + rule_id=f"expire-{storage_expiration_days}-days", + expiration=Expiration(days=365), + ), + ], + ) + await logger.ainfo(f"Setting Minio bucket files to expire in {storage_expiration_days} days", bucket=self.assessment_id) + await self.minio_client.set_bucket_lifecycle(self.assessment_id, config) + await logger.adebug("Uploading to storage", file_path=file_path) file_uuid = uuid.uuid4() await self.minio_client.fput_object(self.assessment_id, f"{file_uuid}", file_path) @@ -56,12 +75,20 @@ class StorageMinio(StorageInterface): raise NotImplementedError async def delete_all_files(self) -> bool: + """ + Deletes all of the files in a bucket. + + For Minio, because we recreate the bucket with the expiration policy on + bucket creation, we want to delete the bucket here as well so the next + upload creates everything correctly. + """ await logger.adebug("Deleting all files from bucket", bucket_name=self.assessment_id) try: files = await self.minio_client.list_objects(self.assessment_id, recursive=True) for file in files: await self.minio_client.remove_object(self.assessment_id, file.object_name) + await self.minio_client.remove_bucket(self.assessment_id) return True except Exception as e: await logger.aexception(e, message="Failed to delete files from bucket", bucket_name=self.assessment_id)