From b648c19bfb71b08535f3fc6ca9fa9530201e5e50 Mon Sep 17 00:00:00 2001
From: Peter Nagy <nagypeterjob@gmail.com>
Date: Wed, 20 Oct 2021 22:29:27 +0200
Subject: [PATCH] feat(es): make elasticsearch plugin work with python3

---
 nixstatsagent/plugins/elasticsearch.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/nixstatsagent/plugins/elasticsearch.py b/nixstatsagent/plugins/elasticsearch.py
index 3437f5f..ff4e896 100644
--- a/nixstatsagent/plugins/elasticsearch.py
+++ b/nixstatsagent/plugins/elasticsearch.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+import urllib.request
 try:
     from urllib.parse import urlparse, urlencode
     from urllib.request import urlopen, Request
@@ -12,7 +13,7 @@
 import plugins
 import json
 import collections
-
+import base64
 
 class Plugin(plugins.BasePlugin):
     __name__ = 'elasticsearch'
@@ -24,6 +25,10 @@ def run(self, config):
         [elasticsearch]
         enabled = yes
         status_page_url = http://127.0.0.1:9200/_stats
+        # In case basic_auth is needed:
+        basic_auth = yes
+        username = <username>
+        password = <password>
         '''
 
         def ascii_encode_dict(data):
@@ -32,8 +37,14 @@ def ascii_encode_dict(data):
 
         results = dict()
         next_cache = dict()
-        request = urllib2.Request(config.get('elasticsearch', 'status_page_url'))
-        raw_response = urllib2.urlopen(request)
+        request = urllib.request.Request(config.get('elasticsearch', 'status_page_url'))
+        basic_auth_enabled = config.get('elasticsearch', 'basic_auth')
+        if bool(basic_auth_enabled):
+            user = config.get('elasticsearch', 'username')
+            password = config.get('elasticsearch', 'password')
+            b64auth = base64.b64encode(bytes("%s:%s" % (user, password), 'ascii'))
+            request.add_header("Authorization", "Basic %s" % b64auth.decode('utf-8'))
+        raw_response = urllib.request.urlopen(request)
         next_cache['ts'] = time.time()
         prev_cache = self.get_agent_cache()  # Get absolute values from previous check
         def flatten(d, parent_key='', sep='_'):
@@ -46,7 +57,7 @@ def flatten(d, parent_key='', sep='_'):
                     items.append((new_key, v))
             return dict(items)
         try:
-            j = flatten(json.loads(raw_response.read(), object_hook=ascii_encode_dict)['_all']['total'])
+            j = flatten(json.loads(raw_response.read())['_all']['total'])
         except Exception:
             return False