Merge pull request #46 from SpecterOps/datalake-purge

Datalake purge and Minion bucket expiration policy

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)

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]:

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
@@ -125,6 +122,7 @@ class NemesisApi(TaskInterface):
     assessment_id: str
     log_level: str
     reprocessing_workers: int
+    storage_expiration_days: int
 
     def __init__(
         self,
@@ -137,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
@@ -147,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")
@@ -159,6 +159,7 @@ class NemesisApi(TaskInterface):
             self.queue_map,
             self.assessment_id,
             self.reprocessing_workers,
+            self.storage_expiration_days,
         )
 
         app.include_router(routes.router)
@@ -202,6 +203,7 @@ class NemesisApiRoutes():
     producers: Dict[NemesisQueue, MessageQueueProducerInterface]
     assessment_id: str
     reprocessing_workers: int
+    storage_expiration_days: int
 
     def __init__(
         self,
@@ -213,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
@@ -223,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"])
@@ -243,7 +247,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 +265,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."""
 
@@ -394,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

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

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)
@@ -55,6 +74,26 @@ class StorageMinio(StorageInterface):
     async def exists(self, file_name: str) -> bool:
         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)
+            raise
+
     async def __aenter__(self):
         return self
 

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"