Add MAC to statistics, to ensure we're talking to the correct process

pyeggsfs/block_service.py

@@ -89,11 +89,17 @@ def deser_write_block(m, rk):
     assert kind == b'w'
     return block_id, crc, size
 
+def deser_stats(m):
+    # prefix=s, length=9
+    kind, token = struct.unpack('<cQ', m)
+    assert kind == b's'
+    return token
+
 def ser_erase_cert(block_id, rk):
     # confirm block_id was erased
     # prefix=E, length=9
     m = struct.pack('<cQ', b'E', block_id)
-    return crypto.add_mac_fixed_len(m, rk)
+    return crypto.add_mac(m, rk)
 
 def ser_fetch_response(size):
     # just return the four byte size followed by the data
@@ -109,11 +115,12 @@ def ser_write_cert(block_id):
     # confirm block was written
     # prefix=W, length=9
     m = struct.pack('<cQ', b'W', block_id)
-    return crypto.add_mac_fixed_len(m, rk)
+    return crypto.add_mac(m, rk)
 
-def ser_stats(used_bytes, free_bytes, used_n, free_n):
-    # prefix=S, length=33
-    return struct.pack('<cQQQQ', b'S', used_bytes, free_bytes, used_n, free_n)
+def ser_stats(token, used_bytes, free_bytes, used_n, free_n, rk):
+    # prefix=S, length=41
+    m = struct.pack('<cQQQQQ', b'S', token, used_bytes, free_bytes, used_n, free_n)
+    return crypto.add_mac(m, rk)
 
 ###########################################
 # main client handling code
@@ -162,10 +169,10 @@ async def handle_client(reader, writer, rk, base_path):
             writer.write(m)
 
         elif kind == b's':
-            # status report (no MAC is required)
-            # TODO we need a MAC so they can verify that they're talking to the correct process
+            # status report (return a MAC so they can verify we are the correct process)
+            token = deser_stats(kind + (await reader.readexactly(8)))
             used_bytes, free_bytes, used_n, free_n = get_stats(base_path)
-            m = ser_stats(used_bytes, free_bytes, used_n, free_n)
+            m = ser_stats(token, used_bytes, free_bytes, used_n, free_n, rk)
             writer.write(m)
 
         elif kind == b'':

pyeggsfs/shuckle.py

@@ -8,7 +8,9 @@
 
 from quart import Quart, request, g, jsonify, render_template
 import asyncio
+import crypto
 import datetime
+import secrets
 import sqlite3
 import struct
 
@@ -18,24 +20,30 @@ app = Quart(__name__)
 # device polling
 ###########################################
 
-async def check_one_device(ip, port):
+async def check_one_device(ip, port, secret_key):
+    rk = crypto.aes_expand_key(bytes.fromhex(secret_key))
     reader, writer = await asyncio.open_connection(ip, port)
-    writer.write(b's')
-    m = await reader.readexactly(33)
-    kind, used_bytes, free_bytes, used_n, free_n = struct.unpack('<cQQQQ', m)
-    assert kind == b'S'
+    token = secrets.token_bytes(8)
+    writer.write(b's' + token)
+    m = await reader.readexactly(49)
+    m = crypto.remove_mac(m, rk)
+    assert m is not None, 'bad mac'
+    kind, got_token, used_bytes, free_bytes, used_n, free_n = struct.unpack('<c8sQQQQ', m)
+    assert kind == b'S', 'bad response kind'
+    assert got_token == token, 'token mismatch'
     return used_bytes, free_bytes, used_n, free_n
 
 async def check_all_devices():
     c = app.db.execute('select id from block_services')
     all_ids = [row[0] for row in c.fetchall()]
     for i in all_ids:
-        c = app.db.execute('select ip, port from block_services where id=? and status!="terminal"', (i,))
+        c = app.db.execute('select ip, port, secret_key from block_services where id=? and status!="terminal"', (i,))
         try:
-            ip, port = c.fetchone()
+            ip, port, secret_key = c.fetchone()
             now = int((datetime.datetime.utcnow() - datetime.datetime(2020, 1, 1)).total_seconds())
-            used_bytes, free_bytes, used_n, free_n = await check_one_device(ip, port)
+            used_bytes, free_bytes, used_n, free_n = await check_one_device(ip, port, secret_key)
         except Exception as e:
+            print(e)
             continue
         app.db.execute('update block_services set last_check=?, used_bytes=?, free_bytes=?, used_n=?, free_n=? where id=?', (now, used_bytes, free_bytes, used_n, free_n, i))