Merge pull request #1412 from p-l-/enh-documentdb

DB: support Amazon DocumentDB

ivre/db/__init__.py

@@ -2108,6 +2108,7 @@ class DBNmap(DBActive):
     backends = {
         "http": ("http", "HttpDBNmap"),
         "mongodb": ("mongo", "MongoDBNmap"),
+        "documentdb": ("document", "DocumentDBNmap"),
         "postgresql": ("sql.postgres", "PostgresDBNmap"),
         "tinydb": ("tiny", "TinyDBNmap"),
     }
@@ -3625,6 +3626,7 @@ class DBView(DBActive):
         "elastic": ("elastic", "ElasticDBView"),
         "http": ("http", "HttpDBView"),
         "mongodb": ("mongo", "MongoDBView"),
+        "documentdb": ("document", "DocumentDBView"),
         "postgresql": ("sql.postgres", "PostgresDBView"),
         "tinydb": ("tiny", "TinyDBView"),
     }
@@ -3793,6 +3795,7 @@ class DBPassive(DB):
     backends = {
         "http": ("http", "HttpDBPassive"),
         "mongodb": ("mongo", "MongoDBPassive"),
+        "documentdb": ("document", "DocumentDBPassive"),
         "postgresql": ("sql.postgres", "PostgresDBPassive"),
         "sqlite": ("sql.sqlite", "SqliteDBPassive"),
         "tinydb": ("tiny", "TinyDBPassive"),
@@ -4268,6 +4271,7 @@ class DBAgent(DB):
 
     backends = {
         "mongodb": ("mongo", "MongoDBAgent"),
+        "documentdb": ("document", "DocumentDBAgent"),
         "tinydb": ("tiny", "TinyDBAgent"),
     }
 
@@ -4722,6 +4726,7 @@ class DBFlow(DB):
 
     backends = {
         "mongodb": ("mongo", "MongoDBFlow"),
+        "documentdb": ("document", "DocumentDBFlow"),
         "postgresql": ("sql.postgres", "PostgresDBFlow"),
         "tinydb": ("tiny", "TinyDBFlow"),
     }

ivre/db/document.py

@@ -0,0 +1,55 @@
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# This file is part of IVRE.
+# Copyright 2011 - 2022 Pierre LALET <pierre@droids-corp.org>
+#
+# IVRE is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# IVRE is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
+# License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with IVRE. If not, see <http://www.gnu.org/licenses/>.
+
+"""This sub-module contains functions to interact with the Amazon
+DocumentDB databases.
+
+"""
+
+
+from ivre.db.mongo import (
+    MongoDBActive,
+    MongoDBNmap,
+    MongoDBView,
+    MongoDBPassive,
+    MongoDBAgent,
+    MongoDBFlow,
+)
+
+
+class DocumentDBNmap(MongoDBNmap):
+    pass
+
+
+class DocumentDBView(MongoDBView):
+    # DocumentDB has no support for text indexes
+    indexes = MongoDBActive.indexes
+    schema_migrations_indexes = MongoDBActive.schema_migrations_indexes
+
+
+class DocumentDBPassive(MongoDBPassive):
+    pass
+
+
+class DocumentDBAgent(MongoDBAgent):
+    pass
+
+
+class DocumentDBFlow(MongoDBFlow):
+    pass

ivre/db/mongo.py

@@ -182,7 +182,14 @@ def db_client(self):
             if self.mongodb_srv is not None:
                 self._db_client = pymongo.MongoClient(self.mongodb_srv)
             else:
-                self._db_client = pymongo.MongoClient(host=self.host, **self.params)
+                self._db_client = pymongo.MongoClient(
+                    host=self.host,
+                    **{
+                        param: value
+                        for param, value in self.params.items()
+                        if not param.startswith("colname_")
+                    },
+                )
             return self._db_client
 
     @property
@@ -842,22 +849,28 @@ class MongoDBActive(MongoDB, DBActive):
                     ("ports.service_product", pymongo.ASCENDING),
                     ("ports.service_version", pymongo.ASCENDING),
                 ],
-                {},
+                {"name": "ivre.hosts.$ports.service"},
             ),
             ([("ports.scripts.id", pymongo.ASCENDING)], {}),
             (
                 [
                     ("ports.scripts.http-headers.name", pymongo.ASCENDING),
                     ("ports.scripts.http-headers.value", pymongo.ASCENDING),
                 ],
-                {"sparse": True},
+                {
+                    "sparse": True,
+                    "name": "ivre.hosts.$ports.$scripts.http-headers",
+                },
             ),
             (
                 [
                     ("ports.scripts.http-app.application", pymongo.ASCENDING),
                     ("ports.scripts.http-app.version", pymongo.ASCENDING),
                 ],
-                {"sparse": True},
+                {
+                    "sparse": True,
+                    "name": "ivre.hosts.$ports.$scripts.http-app",
+                },
             ),
             (
                 [("ports.scripts.dns-domains.parents", pymongo.ASCENDING)],
@@ -993,7 +1006,7 @@ class MongoDBActive(MongoDB, DBActive):
                     ("traces.hops.ipaddr_1", pymongo.ASCENDING),
                     ("traces.hops.ttl", pymongo.ASCENDING),
                 ],
-                {},
+                {"name": "ivre.hosts.$traces.$hops"},
             ),
             (
                 [