notification_proxy: Add heartbeat and startup monitoring

4 files changed, 138 insertions, 38 deletions

diff --git a/notification_proxy.py b/notification_proxy.py
index ae4d73e..d30a23d 100644
--- a/notification_proxy.py
+++ b/notification_proxy.py
@@ -4,86 +4,169 @@ import email.utils
 import hmac
 from email.mime.text import MIMEText
 from datetime import datetime
+import time
 import functools
 import json
 import binascii
+import uwsgidecorators
+
+import sqlite3
 
 from flask import Flask, request, abort
 
 app = Flask(__name__)
 app.config.from_pyfile('config.py')
 
-smtp_server = "smtp.sendgrid.net"
-port = 465
+db = sqlite3.connect(app.config['SQLITE_DB'], check_same_thread=False)
+with db as conn:
+    conn.execute('''CREATE TABLE IF NOT EXISTS seqs_seen
+            (route_name TEXT PRIMARY KEY,
+            seq INTEGER)''')
+    conn.execute('''CREATE TABLE IF NOT EXISTS time_seen
+            (route_name TEXT PRIMARY KEY)''')
+
+    conn.execute('''CREATE TABLE IF NOT EXISTS heartbeats_seen
+            (route_name TEXT PRIMARY KEY,
+            timestamp INTEGER,
+            notified INTEGER)''')
+    # Clear table on startup to avoid spurious notifications
+    conn.execute('''DELETE FROM heartbeats_seen''')
 
 mail_routes = {}
-def mail_route(name, receiver, subject):
+
+def mail_route(name, receiver, subject, secret):
     def wrap(func):
         global routes
-        mail_routes[name] = (receiver, subject, func)
+        mail_routes[name] = (receiver, subject, func, secret)
         return func
     return wrap
 
 
-def authenticate(secret):
-    def wrap(func):
-        func.last_seqnum = 0
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            if not request.is_json:
+def authenticate(route_name, secret, clock_delta_tolerance:'s'=120):
+    with db as conn:
+        if not request.is_json:
+            print('Rejecting notification: Incorrect content type')
+            abort(400)
+
+        if not 'auth' in request.json and 'payload' in request.json:
+            print('Rejecting notification: signature or payload not found')
+            abort(400)
+
+        if not isinstance(request.json['auth'], str):
+            print('Rejecting notification: signature is of incorrect type')
+            abort(400)
+        their_digest = binascii.unhexlify(request.json['auth'])
+
+        our_digest = hmac.digest(secret.encode('utf-8'), request.json['payload'].encode('utf-8'), 'sha256')
+        if not hmac.compare_digest(their_digest, our_digest):
+            print('Rejecting notification: Incorrect signature')
+            abort(403)
+
+        try:
+            payload = json.loads(request.json['payload'])
+        except:
+            print('Rejecting notification: Payload is not JSON')
+            abort(400)
+
+        last_seqnum = conn.execute('SELECT seq FROM seqs_seen WHERE route_name = ?', (route_name,)).fetchone() or 0
+        # We can check for seq here: Only an attacker with knowledge of the secret would be able to remove
+        # seq from a message. This means for a single key, only messages with or without seq may ever be used.
+        if 'seq' in payload:
+            seq = payload['seq']
+            if not isinstance(seq, int):
+                print('Rejecting notification: seq of wrong type')
                 abort(400)
 
-            if not 'auth' in request.json and 'payload' in request.json:
+            if seq <= last_seqnum:
+                print('Rejecting notification: seq out of order')
                 abort(400)
 
-            if not isinstance(request.json['auth'], str):
-                abort(400)
-            their_digest = binascii.unhexlify(request.json['auth'])
+            conn.execute('INSERT OR REPLACE INTO seqs_seen VALUES (?, ?)', (route_name, seq))
 
-            our_digest = hmac.digest(secret.encode('utf-8'), request.json['payload'].encode('utf-8'), 'sha256')
-            if not hmac.compare_digest(their_digest, our_digest):
-                abort(403)
+        elif last_seqnum:
+            print('Rejecting notification: seq not included but past messages included seq')
+            abort(400)
 
-            try:
-                payload = json.loads(request.json['payload'])
-            except:
+        msg_time = None
+        if 'time' in payload:
+            msg_time = payload['time']
+            if not isinstance(msg_time, int):
+                print('Rejecting notification: time of wrong type')
                 abort(400)
 
-            if not isinstance(payload['seq'], int) or payload['seq'] <= func.last_seqnum:
+            if abs(msg_time - int(time.time())) > clock_delta_tolerance:
+                print('Rejecting notification: timestamp too far in the future or past')
                 abort(400)
 
-            func.last_seqnum = payload['seq']
-            del payload['seq']
-            return func(payload)
-        return wrapper
-    return wrap
+            conn.execute('INSERT OR REPLACE INTO time_seen VALUES (?)', (route_name,))
 
-@mail_route('klingel', 'computerstuff@jaseg.de', 'It rang!')
-@authenticate(app.config['SECRET_KLINGEL'])
-def klingel(_):
-    return f'Date: {datetime.utcnow().isoformat()}'
+        elif conn.execute('SELECT * FROM time_seen WHERE route_name = ?', (route_name,)).fetchone():
+            print('Rejecting notification: time not included but past messages included time')
+            abort(400)
 
+        if msg_time is None:
+            msg_time = int(time.time())
 
-@app.route('/notify/<route_name>', methods=['POST'])
-def notify(route_name):
+        return msg_time, payload['scope'], payload['d']
+
+@mail_route('klingel', 'computerstuff@jaseg.de', 'It rang!', app.config['SECRET_KLINGEL'])
+def klingel(rms=None, capture=None, **kwargs):
+    return f'rms={rms}\ncapture={capture}\nextra_args={kwargs}'
+
+
+def send_mail(route_name, receiver, subject, body):
     try:
         context = ssl.create_default_context()
-        smtp = smtplib.SMTP_SSL(smtp_server, port)
+        smtp = smtplib.SMTP_SSL(app.config['SMTP_HOST'], app.config['SMTP_PORT'])
         smtp.login('apikey', app.config['SENDGRID_APIKEY'])
 
         sender = f'{route_name}@{app.config["DOMAIN"]}'
 
-        receiver, subject, func = mail_routes[route_name]
-
-        msg = MIMEText(func() or subject)
+        msg = MIMEText(body)
         msg['Subject'] = subject
         msg['From'] = sender
         msg['To'] = receiver
         msg['Date'] = email.utils.formatdate()
+
         smtp.sendmail(sender, receiver, msg.as_string())
     finally:
         smtp.quit() 
+
+@app.route('/v1/notify/<route_name>', methods=['POST'])
+def notify(route_name):
+    receiver, notify_subject, func, secret = mail_routes[route_name]
+    msg_time, scope, kwargs = authenticate(route_name, secret)
+
+    if scope == 'default':
+        # Exceptions will yield a 500 error
+        body = func(**kwargs)
+        send_mail(route_name, receiver, notify_subject, body or 'empty message')
+
+    elif scope == 'boot':
+        formatted = datetime.utcfromtimestamp(msg_time).isoformat()
+        send_mail(route_name, receiver, 'System startup', f'System powered up at {formatted}')
+
+    elif scope == 'heartbeat':
+        with db as conn:
+            conn.execute('INSERT OR REPLACE INTO heartbeats_seen VALUES (?, ?, 0)', (route_name, int(time.time())))
+
     return 'success'
 
+@uwsgidecorators.timer(60)
+def heartbeat_timer(_uwsgi_signum):
+    threshold = int(time.time()) - app.config['HEARTBEAT_TIMEOUT']
+    with db as conn:
+        for route, ts in db.execute(
+                'SELECT route_name, timestamp FROM heartbeats_seen WHERE timestamp <= ? AND notified == 0',
+                (threshold,)).fetchall():
+            print(f'Heartbeat expired for {route}: {ts} < {threshold}')
+
+            receiver, *_ = mail_routes[route]
+            last = datetime.utcfromtimestamp(ts).isoformat()
+
+            send_mail(route, receiver, 'Heartbeat timeout', f'Last heartbeat at {last}')
+            db.execute('UPDATE heartbeats_seen SET notified = ? WHERE route_name = ?', (int(time.time()), route))
+
 if __name__ == '__main__':
     app.run()
+
diff --git a/notification_proxy_config.py.j2 b/notification_proxy_config.py.j2
index ea53e34..2ecf571 100644
--- a/notification_proxy_config.py.j2
+++ b/notification_proxy_config.py.j2
@@ -1,5 +1,9 @@
 
 SENDGRID_APIKEY = '{{lookup('file', 'notification_proxy_sendgrid_apikey.txt')}}'
 DOMAIN = 'automation.jaseg.de'
+SMTP_HOST = "smtp.sendgrid.net"
+SMTP_PORT = 465
+HEARTBEAT_TIMEOUT = 300
+SQLITE_DB = '{{notification_proxy_sqlite_dbfile}}'
 
 SECRET_KLINGEL = '{{lookup('password', 'notification_proxy_klingel_secret.txt length=32')}}'
diff --git a/playbook.yml b/playbook.yml
index 2db45bc..7c7c95d 100644
--- a/playbook.yml
+++ b/playbook.yml
@@ -12,7 +12,7 @@
 
     - name: Install host requisites
       dnf:
-          name: nginx,uwsgi,python3-flask,python3-flask-wtf,uwsgi-plugin-python3,certbot,python3-certbot-nginx,libselinux-python,git,iptables-services,python3-pycryptodomex,zip
+          name: nginx,uwsgi,python3-flask,python3-flask-wtf,uwsgi-plugin-python3,certbot,python3-certbot-nginx,libselinux-python,git,iptables-services,python3-pycryptodomex,zip,python3-uwsgidecorators
           state: latest
 
     - name: Disable password-based root login
diff --git a/setup_notification_proxy.yml b/setup_notification_proxy.yml
index 3f86412..b47af05 100644
--- a/setup_notification_proxy.yml
+++ b/setup_notification_proxy.yml
@@ -1,4 +1,8 @@
 ---
+- name: Set local facts
+  set_fact:
+      notification_proxy_sqlite_dbfile: /var/lib/notification-proxy/db.sqlite3
+
 - name: Create notification proxy worker user and group
   user:
       name: uwsgi-notification-proxy
@@ -14,7 +18,7 @@
       state: directory
       owner: uwsgi-notification-proxy
       group: uwsgi
-      mode: 0550
+      mode: 0750
 
 - name: Copy webapp sources
   copy:
@@ -46,3 +50,12 @@
       name: uwsgi-app@notification-proxy.socket
       enabled: yes
 
+- name: Create sqlite db file
+  file:
+    path: "{{notification_proxy_sqlite_dbfile}}"
+    owner: uwsgi-notification-proxy
+    group: uwsgi
+    mode: 0660
+    state: touch
+
+