Possibility to protect some endpoints with basic-auth

user:passwd pair should be configured in restapi section of main
configuration file in following format:
restapi:
  auth: 'username:password'

Plus implemented some simple routing mechanisms:
GET /foo => do_GET_foo()
POST /bar => do_POST_bar()

patroni/api.py

@@ -1,3 +1,4 @@
+import base64
 import fcntl
 import json
 import logging
@@ -10,9 +11,36 @@ from threading import Thread
 logger = logging.getLogger(__name__)
 
 
+def check_auth(func):
+    """Decorator function to check authorization header.
+
+    Usage example:
+    @check_auth
+    def do_PUT_foo():
+        pass
+    """
+    def wrapper(handler):
+        if handler.check_auth_header():
+            return func(handler)
+    return wrapper
+
+
 class RestApiHandler(BaseHTTPRequestHandler):
 
+    def send_auth_request(self, body):
+        self.send_response(401)
+        self.send_header('WWW-Authenticate', 'Basic realm=\"Patroni\"')
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+
+    def check_auth_header(self):
+        auth_header = self.headers.get('Authorization')
+        status = self.server.check_auth_header(auth_header)
+        return not status or self.send_auth_request(status)
+
     def do_GET(self):
+        """Default method for processing all GET requests which can not be routed to other methods"""
+
         response = self.get_postgresql_status()
 
         path = '/master' if self.path == '/' else self.path
@@ -23,6 +51,31 @@ class RestApiHandler(BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(json.dumps(response).encode('utf-8'))
 
+    @check_auth
+    def do_GET_sampleauth(self):
+        self.send_response(200)
+        self.send_header('Content-Type', 'text/html')
+        self.end_headers()
+        self.wfile.write(b'Hello!')
+
+    def parse_request(self):
+        """Override parse_request method to enrich basic functionality of `BaseHTTPRequestHandler` class
+
+        Original class can only invoke do_GET, do_POST, do_PUT, etc method implementations if they are defined.
+        But we would like to have at least some simple routing mechanism, i.e.:
+        GET /uri1/part2 request should invoke `do_GET_uri1()`
+        POST /other should invoke `do_POST_other()`
+
+        If the `do_<REQUEST_METHOD>_<first_part_url>` method does not exists we'll fallback to original behavior."""
+
+        ret = BaseHTTPRequestHandler.parse_request(self)
+        if ret:
+            mname = self.path.lstrip('/').split('/')[0]
+            mname = self.command + ('_' + mname if mname else '')
+            if hasattr(self, 'do_' + mname):
+                self.command = mname
+        return ret
+
     def get_postgresql_status(self):
         try:
             row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
@@ -52,6 +105,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
 class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
 
     def __init__(self, patroni, config):
+        self._auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None
         self.connection_string = 'http://{}/patroni'.format(config.get('connect_address', None) or config['listen'])
         host, port = config['listen'].split(':')
         HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
@@ -71,3 +125,13 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
     def _set_fd_cloexec(fd):
         flags = fcntl.fcntl(fd, fcntl.F_GETFD)
         fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
+
+    def check_basic_auth_key(self, key):
+        return self._auth_key == key
+
+    def check_auth_header(self, auth_header):
+        if self._auth_key:
+            if auth_header is None:
+                return 'no auth header received'
+            if not auth_header.startswith('Basic ') or not self.check_basic_auth_key(auth_header[6:]):
+                return 'not authenticated'

postgres0.yml

@@ -4,6 +4,7 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8008
   connect_address: 127.0.0.1:8008
+  auth: 'username:password'
 etcd:
   scope: *scope
   ttl: *ttl

postgres1.yml

@@ -4,6 +4,7 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8009
   connect_address: 127.0.0.1:8009
+  auth: 'username:password'
 etcd:
   scope: *scope
   ttl: *ttl

tests/test_api.py

@@ -3,13 +3,27 @@ import unittest
 
 from patroni.api import RestApiHandler, RestApiServer
 from six import BytesIO as IO
+from six.moves import BaseHTTPServer
 from test_postgresql import psycopg2_connect
 
 
+def nop(*args, **kwargs):
+    pass
+
+
 def throws(*args, **kwargs):
     raise psycopg2.OperationalError()
 
 
+class Mock_BaseServer__is_shut_down:
+
+    def set(self):
+        pass
+
+    def clear(self):
+        pass
+
+
 class MockPostgresql:
 
     def connection(self):
@@ -37,7 +51,7 @@ class MockRequest:
 class MockRestApiServer(RestApiServer):
 
     def __init__(self, Handler, path, *args):
-        self.patroni = MockPatroni()
+        super(MockRestApiServer, self).__init__(MockPatroni(), {'listen': '127.0.0.1:8008', 'auth': 'test:test'})
         if len(args) > 0:
             self.query = args[0]
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
@@ -46,8 +60,20 @@ class MockRestApiServer(RestApiServer):
 class TestRestApiHandler(unittest.TestCase):
 
     def __init__(self, method_name='runTest'):
+        self.setUp = self.set_up
         super(TestRestApiHandler, self).__init__(method_name)
 
+    def set_up(self):
+        BaseHTTPServer.HTTPServer.__init__ = nop
+        RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down()
+        RestApiServer._BaseServer__shutdown_request = True
+        RestApiServer.socket = 0
+
     def test_do_GET(self):
         MockRestApiServer(RestApiHandler, b'GET /')
         MockRestApiServer(RestApiHandler, b'GET /', throws)
+
+    def test_do_GET_sampleauth(self):
+        MockRestApiServer(RestApiHandler, b'GET /sampleauth')
+        MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization:')
+        MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization: Basic dGVzdDp0ZXN0')

tests/test_patroni.py

@@ -14,6 +14,7 @@ from patroni.etcd import Etcd
 from patroni import Patroni, main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
+from test_api import Mock_BaseServer__is_shut_down
 from test_etcd import Client, etcd_read, etcd_write
 from test_ha import true, false
 from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
@@ -32,15 +33,6 @@ def time_sleep(*args):
     raise SleepException()
 
 
-class Mock_BaseServer__is_shut_down:
-
-    def set(self):
-        pass
-
-    def clear(self):
-        pass
-
-
 class TestPatroni(unittest.TestCase):
 
     def __init__(self, method_name='runTest'):