Fix issues with elasticsearch API

mmguero-dev · Oct 24, 2023 · 5ae09a7 · 5ae09a7
1 parent 81c3ab7
commit 5ae09a7
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 58 deletions.
diff --git a/api/project/__init__.py b/api/project/__init__.py
@@ -166,20 +166,7 @@
 
 opensearchUrl = app.config["OPENSEARCH_URL"]
 dashboardsUrl = app.config["DASHBOARDS_URL"]
-
 databaseMode = malcolm_utils.DatabaseModeStrToEnum(app.config["OPENSEARCH_PRIMARY"])
-if databaseMode == malcolm_utils.DatabaseMode.ElasticsearchRemote:
-    import elasticsearch as DatabaseImport
-    from elasticsearch_dsl import Search as SearchClass
-else:
-    import opensearchpy as DatabaseImport
-    from opensearchpy import Search as SearchClass
-
-DatabaseClass = (
-    DatabaseImport.Elasticsearch
-    if databaseMode == malcolm_utils.DatabaseMode.ElasticsearchRemote
-    else DatabaseImport.OpenSearch
-)
 
 opensearchLocal = (databaseMode == malcolm_utils.DatabaseMode.OpenSearchLocal) or (
     opensearchUrl == 'http://opensearch:9200'
@@ -190,26 +177,39 @@
     if (not opensearchLocal)
     else defaultdict(lambda: None)
 )
+
+DatabaseInitArgs = {}
+if urlparse(opensearchUrl).scheme == 'https':
+    DatabaseInitArgs['verify_certs'] = opensearchSslVerify
+    DatabaseInitArgs['ssl_assert_hostname'] = False
+    DatabaseInitArgs['ssl_show_warn'] = False
+
 if opensearchCreds['user'] is not None:
-    opensearchHttpAuth = f"{opensearchCreds['user']}:{opensearchCreds['password']}"
+    opensearchHttpAuth = (opensearchCreds['user'], opensearchCreds['password'])
     opensearchReqHttpAuth = HTTPBasicAuth(opensearchCreds['user'], opensearchCreds['password'])
 else:
     opensearchHttpAuth = None
     opensearchReqHttpAuth = None
 
-if urlparse(opensearchUrl).scheme == 'https':
-    databaseClient = DatabaseClass(
-        hosts=[opensearchUrl],
-        http_auth=opensearchHttpAuth,
-        verify_certs=opensearchSslVerify,
-        ssl_assert_hostname=False,
-        ssl_show_warn=False,
-    )
+if databaseMode == malcolm_utils.DatabaseMode.ElasticsearchRemote:
+    import elasticsearch as DatabaseImport
+    from elasticsearch_dsl import Search as SearchClass
+
+    DatabaseClass = DatabaseImport.Elasticsearch
+    if opensearchHttpAuth:
+        DatabaseInitArgs['basic_auth'] = opensearchHttpAuth
 else:
-    databaseClient = DatabaseClass(
-        hosts=[opensearchUrl],
-        http_auth=opensearchHttpAuth,
-    )
+    import opensearchpy as DatabaseImport
+    from opensearchpy import Search as SearchClass
+
+    DatabaseClass = DatabaseImport.OpenSearch
+    if opensearchHttpAuth:
+        DatabaseInitArgs['http_auth'] = opensearchHttpAuth
+
+databaseClient = DatabaseClass(
+    hosts=[opensearchUrl],
+    **DatabaseInitArgs,
+)
 
 
 def deep_get(d, keys, default=None):

diff --git a/shared/bin/pcap_watcher.py b/shared/bin/pcap_watcher.py
@@ -60,23 +60,29 @@
 ###################################################################################################
 pdbFlagged = False
 args = None
-opensearchHttpAuth = None
 scriptName = os.path.basename(__file__)
 scriptPath = os.path.dirname(os.path.realpath(__file__))
 origPath = os.getcwd()
 shuttingDown = [False]
 DEFAULT_NODE_NAME = os.getenv('PCAP_NODE_NAME', 'malcolm')
-DatabaseClass, SearchClass, ConnectionError, ConnectionTimeout, AuthenticationException = None, None, None, None, None
+DatabaseClass, DatabaseInitArgs, SearchClass, ConnectionError, ConnectionTimeout, AuthenticationException = (
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+)
 
 
 ###################################################################################################
 # watch files written to and moved to this directory
 class EventWatcher:
     def __init__(self, logger=None):
         global args
-        global opensearchHttpAuth
         global shuttingDown
         global DatabaseClass
+        global DatabaseInitArgs
         global SearchClass
         global ConnectionError
         global ConnectionTimeout
@@ -98,23 +104,10 @@ def __init__(self, logger=None):
                 try:
                     try:
                         self.logger.info(f"{scriptName}:\tconnecting to {args.opensearchMode} {args.opensearchUrl}...")
-
-                        if urlparse(args.opensearchUrl).scheme == 'https':
-                            self.openSearchClient = DatabaseClass(
-                                hosts=[args.opensearchUrl],
-                                http_auth=opensearchHttpAuth,
-                                verify_certs=args.opensearchSslVerify,
-                                ssl_assert_hostname=False,
-                                ssl_show_warn=False,
-                                request_timeout=1,
-                            )
-                        else:
-                            self.openSearchClient = DatabaseClass(
-                                hosts=[args.opensearchUrl],
-                                http_auth=opensearchHttpAuth,
-                                request_timeout=1,
-                            )
-
+                        self.openSearchClient = DatabaseClass(
+                            hosts=[args.opensearchUrl],
+                            **DatabaseInitArgs,
+                        )
                         self.logger.debug(f"{scriptName}:\t{self.openSearchClient.cluster.health()}")
 
                         self.openSearchClient.cluster.health(
@@ -274,10 +267,10 @@ def pdb_handler(sig, frame):
 # main
 def main():
     global args
-    global opensearchHttpAuth
     global pdbFlagged
     global shuttingDown
     global DatabaseClass
+    global DatabaseInitArgs
     global SearchClass
     global ConnectionError
     global ConnectionTimeout
@@ -448,14 +441,6 @@ def main():
     if args.verbose > logging.DEBUG:
         sys.tracebacklimit = 0
 
-    if args.opensearchMode == malcolm_utils.DatabaseMode.ElasticsearchRemote:
-        from elasticsearch import Elasticsearch as DatabaseClass
-        from elasticsearch_dsl import Search as SearchClass
-        from elasticsearch.exceptions import ConnectionError, ConnectionTimeout, AuthenticationException
-    else:
-        from opensearchpy import OpenSearch as DatabaseClass, Search as SearchClass
-        from opensearchpy.exceptions import ConnectionError, ConnectionTimeout, AuthenticationException
-
     opensearchIsLocal = (args.opensearchMode == malcolm_utils.DatabaseMode.OpenSearchLocal) or (
         args.opensearchUrl == 'http://opensearch:9200'
     )
@@ -466,9 +451,27 @@ def main():
             args.opensearchUrl = 'http://opensearch:9200'
         elif 'url' in opensearchCreds:
             args.opensearchUrl = opensearchCreds['url']
-    opensearchHttpAuth = (
-        f"{opensearchCreds['user']}:{opensearchCreds['password']}" if opensearchCreds['user'] is not None else None
-    )
+
+    DatabaseInitArgs = {}
+    if args.opensearchMode == malcolm_utils.DatabaseMode.ElasticsearchRemote:
+        from elasticsearch import Elasticsearch as DatabaseClass
+        from elasticsearch_dsl import Search as SearchClass
+        from elasticsearch.exceptions import ConnectionError, ConnectionTimeout, AuthenticationException
+
+        if opensearchCreds['user'] is not None:
+            DatabaseInitArgs['basic_auth'] = (opensearchCreds['user'], opensearchCreds['password'])
+    else:
+        from opensearchpy import OpenSearch as DatabaseClass, Search as SearchClass
+        from opensearchpy.exceptions import ConnectionError, ConnectionTimeout, AuthenticationException
+
+        if opensearchCreds['user'] is not None:
+            DatabaseInitArgs['http_auth'] = (opensearchCreds['user'], opensearchCreds['password'])
+
+    DatabaseInitArgs['request_timeout'] = 1
+    if urlparse(args.opensearchUrl).scheme == 'https':
+        DatabaseInitArgs['verify_certs'] = args.opensearchSslVerify
+        DatabaseInitArgs['ssl_assert_hostname'] = False
+        DatabaseInitArgs['ssl_show_warn'] = False
 
     # handle sigint and sigterm for graceful shutdown
     signal.signal(signal.SIGINT, shutdown_handler)