implement server for nginx auth_request

5 files changed, 187 insertions, 89 deletions

diff --git a/README.md b/README.md
index be8421c..08ce4e4 100644
--- a/README.md
+++ b/README.md
@@ -40,3 +40,21 @@ instance at a time.
 - jid: JID of the bot who sends the tokens to the users.
 - jid\_pw: password of the bot.
 
+
+
+# nginx
+
+## configuration
+
+        location /grafana {
+            auth_request /_auth;
+            # ...
+        }
+
+        location = /_auth {
+            proxy_pass http://localhost:8081/;
+            proxy_pass_request_body off;
+            proxy_set_header Content-Length "";
+            proxy_set_header X-Original-URI "$scheme://$host$request_uri";
+        }
+
diff --git a/functions.py b/functions.py
new file mode 100644
index 0000000..e26e43f
--- /dev/null
+++ b/functions.py
@@ -0,0 +1,80 @@
+import os
+import re
+import time
+import struct
+import hashlib
+from urllib.parse import quote as urlencode
+
+
+def _normalize_token(token):
+    return re.sub(r"[^A-F0-9]", "", token.upper())
+
+
+def _generate_token(username, secret, time):
+    input = "{}{}{}".format(secret, username, time).encode('utf-8')
+    output = struct.unpack(b"<L", hashlib.md5(input).digest()[:4])[0]
+    token = "{:02X}-{:02X}-{:02X}".format(
+        (output >> 16) & 0xff, (output >> 8) & 0xff, output & 0xff)
+    return token
+
+
+def file_lock(lock_file):
+    from contextlib import contextmanager
+
+    @contextmanager
+    def file_lock():
+        try:
+            with open(lock_file, "x") as fh:
+                try:
+                    yield
+                except:
+                    raise
+                finally:
+                    fh.close()
+                    os.remove(lock_file)
+        except FileExistsError:
+            raise Exception("Locking failed on {}".format(lock_file))
+
+    return file_lock()
+
+
+def token_message(username, secret, validsec, url):
+    time_now = int(time.time())
+    time_now_start = int(time_now - time_now % validsec)
+    time_next_end = time_now_start + 2 * validsec
+    token = _generate_token(username, secret, time_now_start)
+    message = "Username: {} Token: {}".format(username, token)
+    message += "\nValid from: {} to: {}".format(
+        time.strftime("%c %Z(%z)", time.gmtime(time_now_start)),
+        time.strftime("%c %Z(%z)", time.gmtime(time_next_end)))
+    if url is not None:
+        message += re.sub('(https?://)(.*)',
+                          ' \\1' + urlencode(username) + ':' + urlencode(token) + '@\\2',
+                          url)
+    return message
+
+
+def send_message(jid, password, recipient, message):
+    import sleekxmpp
+
+    def start(event):
+        cl.send_message(mto=recipient, mtype='chat', mbody=message)
+        cl.disconnect(wait=True)
+
+    cl = sleekxmpp.ClientXMPP(jid, password)
+    cl.add_event_handler("session_start", start, threaded=True)
+    if cl.connect():
+        cl.process(block=True)
+    else:
+        raise Exception("Unable to connect to xmpp server")
+
+
+def verify_token(username, password, conf_secret, conf_validsec):
+    time_now = int(time.time())
+    time_now_start = int(time_now - time_now % conf_validsec)
+    time_prev_start = time_now_start - conf_validsec
+    valid_tokens = list(map(_normalize_token, (
+        _generate_token(username, conf_secret, time_now_start),
+        _generate_token(username, conf_secret, time_prev_start)
+    )))
+    return _normalize_token(password) in valid_tokens
diff --git a/login.py b/login.py
index 019cbf7..24f62a1 100755
--- a/login.py
+++ b/login.py
@@ -1,78 +1,7 @@
 #!/usr/bin/env python3.4
 import os
-import re
 import sys
-import time
-import struct
-import hashlib
-from urllib.parse import quote as urlencode
-
-# To speed up start time load some modules only as needed
-
-if sys.version_info < (3, 0):
-    raise Exception("Require python3+")
-
-
-def file_lock(lock_file):
-    from contextlib import contextmanager
-
-    @contextmanager
-    def file_lock():
-        try:
-            with open(lock_file, "x") as fh:
-                try:
-                    yield
-                except:
-                    raise
-                finally:
-                    fh.close()
-                    os.remove(lock_file)
-        except FileExistsError:
-            raise Exception("Locking failed on {}".format(lock_file))
-    return file_lock()
-
-
-def send_message(jid, password, recipient, message):
-    import sleekxmpp
-
-    def start(event):
-        cl.send_message(mto=recipient, mtype='chat', mbody=message)
-        cl.disconnect(wait=True)
-
-    cl = sleekxmpp.ClientXMPP(jid, password)
-    cl.add_event_handler("session_start", start, threaded=True)
-    if cl.connect():
-        cl.process(block=True)
-    else:
-        raise Exception("Unable to connect to xmpp server")
-
-
-def generate_token(username, secret, time):
-    input = "{}{}{}".format(secret, username, time).encode('utf-8')
-    output = struct.unpack(b"<L", hashlib.md5(input).digest()[:4])[0]
-    token = "{:02X}-{:02X}-{:02X}".format(
-            (output >> 16) & 0xff, (output >> 8) & 0xff, output & 0xff)
-    return token
-
-
-def token_message(username, secret, validsec):
-    time_now = int(time.time())
-    time_now_start = int(time_now - time_now % validsec)
-    time_next_end = time_now_start + 2 * validsec
-    token = generate_token(username, secret, time_now_start)
-    message = "Username: {} Token: {}".format(username, token)
-    message += "\nValid from: {} to: {}".format(
-        time.strftime("%c %Z(%z)", time.gmtime(time_now_start)),
-        time.strftime("%c %Z(%z)", time.gmtime(time_next_end)))
-    message += "\nRequested by: {} for: {} on: {}".format(
-        os.getenv("IP"), ascii(os.getenv("URI")), os.getenv("HTTP_HOST"))
-    message += "\nhttps://{}:{}@{}{}".format(urlencode(username),
-        token, os.getenv("HTTP_HOST"), urlencode(os.getenv("URI")))
-    return message
-
-
-def normalize_token(token):
-    return re.sub(r"[^A-F0-9]", "", token.upper())
+import functions
 
 
 def run(config):
@@ -89,25 +18,20 @@ def run(config):
     if password == "" and username in conf_users:
         # avoid spamming by allowing only one message sent at a time
         lockfile = os.path.basename(__file__)
-        with file_lock("/tmp/lock." + lockfile):
-            message = token_message(username, conf_secret, conf_validsec)
+        with functions.file_lock("/tmp/lock." + lockfile):
+            message = functions.token_message(username, conf_secret, conf_validsec,
+                                              os.getenv("URI"), os.getenv("HTTP_HOST"))
             if os.getenv("SKIP_XMPP"):  # used for testing
                 print(message)
             else:
-                send_message(conf_jid, conf_jid_pw, username, message)
+                functions.send_message(conf_jid, conf_jid_pw, username, message)
     elif username in conf_users:
-        time_now = int(time.time())
-        time_now_start = int(time_now - time_now % conf_validsec)
-        time_prev_start = time_now_start - conf_validsec
-        valid_tokens = list(map(normalize_token, (
-            generate_token(username, conf_secret, time_now_start),
-            generate_token(username, conf_secret, time_prev_start)
-        )))
-        if normalize_token(password) in valid_tokens:
-            return os.EX_OK  # grant access
+        if functions.verify_token(username, password, conf_secret, conf_validsec):
+            return os.EX_OK
 
     return os.EX_NOPERM  # fail by default
 
+
 if __name__ == "__main__":
     config = dict(map(lambda kv: kv.split("="),
                       os.getenv("CONTEXT").split(";")))
diff --git a/login_test.py b/login_test.py
index 4557de8..2bdbef3 100755
--- a/login_test.py
+++ b/login_test.py
@@ -1,17 +1,17 @@
 #!/usr/bin/env python3
 import unittest
-import login
+import functions
 
 
 class TestStringMethods(unittest.TestCase):
     def test_normalize(self):
-        self.assertEqual(login.normalize_token("A4-B4-C5"),
+        self.assertEqual(functions._normalize_token("A4-B4-C5"),
             "A4B4C5")
-        self.assertEqual(login.normalize_token("a4-b4-c5"),
+        self.assertEqual(functions._normalize_token("a4-b4-c5"),
             "A4B4C5")
-        self.assertEqual(login.normalize_token("a4b4c5"),
+        self.assertEqual(functions._normalize_token("a4b4c5"),
             "A4B4C5")
-        self.assertEqual(login.normalize_token("A4B4C5"),
+        self.assertEqual(functions._normalize_token("A4B4C5"),
             "A4B4C5")
 
 
diff --git a/server.py b/server.py
new file mode 100755
index 0000000..4b72f83
--- /dev/null
+++ b/server.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+import time
+import binascii
+import random
+import argparse
+import functions
+import logging
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+logging.basicConfig(level=logging.INFO)
+
+LAST_REQUEST_TIME = 0
+CACHE = {}
+
+
+def send_token(conf, username, orig_uri):
+    message = functions.token_message(username, conf.secret, conf.validsec, orig_uri)
+    if conf.skip_xmpp:  # used for testing
+        print(message)
+    else:
+        functions.send_message(conf.jid, conf.password, username, message)
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        global LAST_REQUEST_TIME, CACHE
+        if 'Authorization' in self.headers:
+            method, value = self.headers['Authorization'].split(' ')
+            if method != 'Basic':
+                self.send_response(400, 'Unsupported authentication method')
+            elif value in CACHE and CACHE[value] > time.time() - 60:  # cache cred for 60s for performance
+                logging.info("Authorized (cached) %s", value)
+                self.send_response(200, "OK go forward")
+            else:
+                username, password = binascii.a2b_base64(value.encode('utf-8')).decode('utf-8').split(':')
+                if password == "" and username in conf.users:
+                    if LAST_REQUEST_TIME == 0 or time.time() - LAST_REQUEST_TIME > 15:  # max 1 msg per 15 sec
+                        LAST_REQUEST_TIME = time.time()
+                        send_token(conf, username, self.headers['X-Original-URI'])
+                        self.send_response(401, "Token sent, retry")
+                    else:
+                        self.send_response(429, 'Too Many Requests')
+                else:
+                    if functions.verify_token(username, password, conf.secret, conf.validsec):
+                        logging.info("Authorized %s", username)
+                        CACHE[value] = time.time()
+                        self.send_response(200, "OK go forward")
+                    else:
+                        logging.info("Denied %s", username)
+                        self.send_response(401, "Authentication failed, username or password wrong")
+        else:
+            self.send_response(401)
+            self.send_header("WWW-Authenticate", "Basic realm=\"xmppmessage auth\"")
+
+        self.end_headers()
+
+
+def run(conf):
+    httpd = HTTPServer((conf.server_host, conf.server_port), RequestHandler)
+    httpd.conf = conf
+    httpd.serve_forever()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--secret', default="".join([chr(random.randint(ord('0'), ord('Z'))) for x in range(20)]))
+    parser.add_argument('--validsec', type=int, default=60 * 60 * 48)
+    parser.add_argument('--user', '-u', nargs='+', default=['yvesf@xapek.org', 'marc@xapek.org'], dest='users')
+    parser.add_argument('--jid', help="Bot jid", default="bot@xapek.org")
+    parser.add_argument('--password', help="Bot jid password")
+    parser.add_argument('--server-host', default="127.0.0.1")
+    parser.add_argument('--server-port', default=8081, type=int)
+    parser.add_argument('--skip-xmpp', default=False, type=bool)
+
+    conf = parser.parse_args()
+    run(conf)