Add mongodb authentication, encryption, and extended flags

mcrit/config/QueueConfig.py

@@ -13,6 +13,9 @@ class QueueConfig(ConfigInterface):
     # By default, MongoDbStorage's DB's name and MongoQueue's DB's name are both "mcrit"
     # Changing one DB name here or at runtime DOES NOT change the other name!
     QUEUE_MONGODB_DBNAME: str = "mcrit"
+    QUEUE_MONGODB_USERNAME: str = ""
+    QUEUE_MONGODB_PASSWORD: str = ""
+    QUEUE_MONGODB_FLAGS: str = ""
     QUEUE_MONGODB_COLLECTION_NAME: str = "queue"
     QUEUE_PORT: ... = None
     QUEUE_TIMEOUT: int = 300

mcrit/config/StorageConfig.py

@@ -15,6 +15,9 @@ class StorageConfig(ConfigInterface):
     # By default, MongoDbStorage's DB's name and MongoQueue's DB's name are both "mcrit"
     # Changing one DB name here or at runtime DOES NOT change the other name!
     STORAGE_MONGODB_DBNAME: str = "mcrit"
+    STORAGE_MONGODB_USERNAME: str = None
+    STORAGE_MONGODB_PASSWORD: str = None
+    STORAGE_MONGODB_FLAGS: str = ""
     STORAGE_PORT: ... = None
     # Once MinHashes have been calculated, discard disassembly from function entries
     STORAGE_DROP_DISASSEMBLY: bool = False

mcrit/libs/mongoqueue.py

@@ -70,7 +70,13 @@ def __init__(self, queue_config, consumer_id, timeout=300, max_attempts=3):
     def _getCollection(self):
         # because of gunicorn and forking workers, we want to delay creation of MongoClient until actual usage and avoid it within __init__()
         if self.collection is None:
-            db = MongoClient(host=self.queue_config.QUEUE_SERVER, port=self.queue_config.QUEUE_PORT, connect=False)
+            userpw_url = f"{self.queue_config.QUEUE_MONGODB_USERNAME}:{self.queue_config.QUEUE_MONGODB_PASSWORD}@" if self.queue_config.QUEUE_MONGODB_USERNAME and len(self.queue_config.QUEUE_MONGODB_USERNAME) > 0 and self.queue_config.QUEUE_MONGODB_PASSWORD and len(self.queue_config.QUEUE_MONGODB_PASSWORD) > 0 else ""
+            port_url = f":{self.queue_config.QUEUE_PORT}" if self.queue_config.QUEUE_PORT else ""
+            flags_url = f"?{self.queue_config.QUEUE_MONGODB_FLAGS}" if self.queue_config.QUEUE_MONGODB_FLAGS and len(self.queue_config.QUEUE_MONGODB_FLAGS) > 0 else ""
+
+            mongo_uri = f"mongodb://{userpw_url}{self.queue_config.QUEUE_SERVER}{port_url}/{self.queue_config.QUEUE_MONGODB_DBNAME}{flags_url}"
+
+            db = MongoClient(mongo_uri, connect=False)
             self.collection = db[self.queue_config.QUEUE_MONGODB_DBNAME][self.queue_config.QUEUE_MONGODB_COLLECTION_NAME]
             self.fs = gridfs.GridFS(self.collection.database)
             self.fs_files = self.collection.database["fs.files"]

mcrit/storage/MongoDbStorage.py

@@ -117,11 +117,22 @@ def __init__(self, config: "StorageConfig") -> None:
     def _getDb(self):
         # because of gunicorn and forking workers, we want to delay creation of MongoClient until actual usage and avoid it within __init__()
         if self._database is None:
-            self._initDb(self._storage_config.STORAGE_SERVER, self._storage_config.STORAGE_PORT, self._storage_config.STORAGE_MONGODB_DBNAME)
+            self._initDb(self._storage_config.STORAGE_SERVER, 
+                         self._storage_config.STORAGE_PORT, 
+                         self._storage_config.STORAGE_MONGODB_DBNAME, 
+                         self._storage_config.STORAGE_MONGODB_USERNAME, 
+                         self._storage_config.STORAGE_MONGODB_PASSWORD, 
+                         self._storage_config.STORAGE_MONGODB_FLAGS)
         return self._database
 
-    def _initDb(self, server, port, db_name):
-        self._database = MongoClient(server, port=port, connect=False)[db_name]
+    def _initDb(self, server, port, db_name, username="", password="", flags=""):
+        userpw_url = f"{username}:{password}@" if username and len(username) > 0 and password and len(password) > 0 else ""
+        port_url = f":{port}" if port and len(port) > 0 else ""
+        flags_url = f"?{flags}" if flags and len(flags) > 0 else ""
+
+        mongo_uri = f"mongodb://{userpw_url}{server}{port_url}/{db_name}{flags_url}"
+
+        self._database = MongoClient(mongo_uri, connect=False)[db_name]
         self._ensureIndexAndUnknownFamily()
 
     def _ensureIndexAndUnknownFamily(self) -> None: