From 54eab0aa352b432dcb1c327183156a0187536650 Mon Sep 17 00:00:00 2001
From: Shyam Sunder <sgsunder1@gmail.com>
Date: Sun, 15 Sep 2019 12:45:17 -0400
Subject: [PATCH] server/image-hash: optionally allow for elasticsearch
 authentication

---
 docker-compose.yml                   |  4 ++++
 server/config.yaml.dist              |  2 ++
 server/requirements.txt              |  1 +
 server/szurubooru/config.py          |  4 +++-
 server/szurubooru/func/image_hash.py | 21 ++++++++++++++-------
 server/wait-for-es                   |  9 +++------
 6 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index 345e531c..83faa306 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -24,6 +24,10 @@ services:
       #POSTGRES_PORT: 5432
       #ESEARCH_PORT: 9200
       #ESEARCH_INDEX: szurubooru
+      #ESEARCH_PASSWORD: (empty by default, set if you are using an external
+      #                  source for elasticsearch and want to use HTTP basic
+      #                  authentication for security)
+      #ESEARCH_USER: szurubooru (only used if password is set)
       #LOG_SQL: 0 (1 for verbose SQL logs)
     volumes:
       - "${MOUNT_DATA}:/data"
diff --git a/server/config.yaml.dist b/server/config.yaml.dist
index ff28c6de..465e11a6 100644
--- a/server/config.yaml.dist
+++ b/server/config.yaml.dist
@@ -146,3 +146,5 @@ privileges:
 #    host: localhost
 #    port: 9200
 #    index: szurubooru
+#    user: szurubooru
+#    pass:
diff --git a/server/requirements.txt b/server/requirements.txt
index a142493b..3a312bba 100644
--- a/server/requirements.txt
+++ b/server/requirements.txt
@@ -5,6 +5,7 @@ SQLAlchemy>=1.0.12
 coloredlogs==5.0
 elasticsearch>=5.0.0,<7.0.0
 elasticsearch-dsl>=5.0.0,<7.0.0
+certifi>=2017.11.5
 numpy>=1.8.2
 pillow>=4.3.0
 pynacl==1.2.1
diff --git a/server/szurubooru/config.py b/server/szurubooru/config.py
index 4e174d91..6b531dc5 100644
--- a/server/szurubooru/config.py
+++ b/server/szurubooru/config.py
@@ -40,7 +40,9 @@ def _docker_config() -> Dict:
         'elasticsearch': {
             'host': os.getenv('ESEARCH_HOST'),
             'port': int(os.getenv('ESEARCH_PORT', 9200)),
-            'index': os.getenv('ESEARCH_INDEX', 'szurubooru')
+            'index': os.getenv('ESEARCH_INDEX', 'szurubooru'),
+            'user': os.getenv('ESEARCH_USER', os.getenv('ESEARCH_INDEX', 'szurubooru')),
+            'pass': os.getenv('ESEARCH_PASSWORD', False)
         }
     }
 
diff --git a/server/szurubooru/func/image_hash.py b/server/szurubooru/func/image_hash.py
index b89b2187..e5ae6a37 100644
--- a/server/szurubooru/func/image_hash.py
+++ b/server/szurubooru/func/image_hash.py
@@ -31,11 +31,18 @@ Window = Tuple[Tuple[float, float], Tuple[float, float]]
 NpMatrix = Any
 
 
-def _get_session() -> elasticsearch.Elasticsearch:
+def get_session() -> elasticsearch.Elasticsearch:
+    extra_args = {}
+    if config.config['elasticsearch']['pass']:
+        extra_args['http_auth'] = (
+            config.config['elasticsearch']['user'],
+            config.config['elasticsearch']['pass'])
+        extra_args['scheme'] = 'https'
+        extra_args['port'] = 443
     return elasticsearch.Elasticsearch([{
         'host': config.config['elasticsearch']['host'],
         'port': config.config['elasticsearch']['port'],
-    }])
+    }], **extra_args)
 
 
 def _preprocess_image(content: bytes) -> NpMatrix:
@@ -271,7 +278,7 @@ def add_image(path: str, image_content: bytes) -> None:
     for i in range(MAX_WORDS):
         record['simple_word_' + str(i)] = words[i].tolist()
 
-    _get_session().index(
+    get_session().index(
         index=config.config['elasticsearch']['index'],
         doc_type=ES_DOC_TYPE,
         body=record,
@@ -281,7 +288,7 @@ def add_image(path: str, image_content: bytes) -> None:
 @_safety_blanket(lambda: None)
 def delete_image(path: str) -> None:
     assert path
-    _get_session().delete_by_query(
+    get_session().delete_by_query(
         index=config.config['elasticsearch']['index'],
         doc_type=ES_DOC_TYPE,
         body={'query': {'term': {'path': path}}})
@@ -292,7 +299,7 @@ def search_by_image(image_content: bytes) -> List[Lookalike]:
     signature = _generate_signature(image_content)
     words = _get_words(signature, k=SAMPLE_WORDS, n=MAX_WORDS)
 
-    res = _get_session().search(
+    res = get_session().search(
         index=config.config['elasticsearch']['index'],
         doc_type=ES_DOC_TYPE,
         body={
@@ -333,7 +340,7 @@ def search_by_image(image_content: bytes) -> List[Lookalike]:
 
 @_safety_blanket(lambda: None)
 def purge() -> None:
-    _get_session().delete_by_query(
+    get_session().delete_by_query(
         index=config.config['elasticsearch']['index'],
         doc_type=ES_DOC_TYPE,
         body={'query': {'match_all': {}}},
@@ -344,7 +351,7 @@ def purge() -> None:
 def get_all_paths() -> Set[str]:
     search = (
         elasticsearch_dsl.Search(
-            using=_get_session(),
+            using=get_session(),
             index=config.config['elasticsearch']['index'],
             doc_type=ES_DOC_TYPE)
         .source(['path']))
diff --git a/server/wait-for-es b/server/wait-for-es
index d7229296..b21ea6eb 100755
--- a/server/wait-for-es
+++ b/server/wait-for-es
@@ -6,17 +6,14 @@ Docker helper script. Blocks until the ElasticSearch service is ready.
 
 import logging
 import time
-import elasticsearch
-from szurubooru import config, errors
+from szurubooru import errors
+from szurubooru.func.image_hash import get_session
 
 
 def main():
     print('Looking for ElasticSearch connection...')
     logging.basicConfig(level=logging.ERROR)
-    es = elasticsearch.Elasticsearch([{
-        'host': config.config['elasticsearch']['host'],
-        'port': config.config['elasticsearch']['port'],
-    }])
+    es = get_session()
 
     TIMEOUT = 30
     DELAY = 0.1