server/image-hash: optionally allow for elasticsearch authentication

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"

server/config.yaml.dist

@@ -146,3 +146,5 @@ privileges:
 #    host: localhost
 #    port: 9200
 #    index: szurubooru
+#    user: szurubooru
+#    pass:

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

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)
         }
     }
 

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']))

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 errors
-from szurubooru import config, errors
+from szurubooru.func.image_hash import get_session
 
 
 def main():
     print('Looking for ElasticSearch connection...')
     logging.basicConfig(level=logging.ERROR)
-    es = elasticsearch.Elasticsearch([{
+    es = get_session()
-        'host': config.config['elasticsearch']['host'],
-        'port': config.config['elasticsearch']['port'],
-    }])
 
     TIMEOUT = 30
     DELAY = 0.1