From ee2e55e1747bd74d8454f2727b073681603bc9b0 Mon Sep 17 00:00:00 2001
From: sqozz <sqozz@geekify.de>
Date: Sat, 17 Nov 2018 23:23:33 +0100
Subject: [PATCH] Add config capabilities

---
 config.json                          | 25 +++++++++
 firmwares/ota_test/latest/latest.bin |  0
 update_server.py                     | 81 +++++++++++++++++++++-------
 3 files changed, 88 insertions(+), 18 deletions(-)
 create mode 100644 config.json
 create mode 100644 firmwares/ota_test/latest/latest.bin

diff --git a/config.json b/config.json
new file mode 100644
index 0000000..5461eb8
--- /dev/null
+++ b/config.json
@@ -0,0 +1,25 @@
+{
+	"devices" : {
+		"84:F3:EB:B4:4D:A7" : {
+			"role" : "ota_test",
+			"version" : "latest"
+		},
+		"AA:BB:CC:DD:EE:FF" : {
+			"role" : "non_existant",
+			"version" : "latest"
+		}
+	},
+
+	"firmwares" : {
+		"ota_test" : {
+			"latest" : {
+				"filename" : "latest.bin",
+				"hash" : "d41d8cd98f00b204e9800998ecf8427e"
+			},
+			"backup" : {
+				"filename" : "backup.bin",
+				"hash" : ""
+			}
+		}
+	}
+}
diff --git a/firmwares/ota_test/latest/latest.bin b/firmwares/ota_test/latest/latest.bin
new file mode 100644
index 0000000..e69de29
diff --git a/update_server.py b/update_server.py
index bbfafce..01e285d 100644
--- a/update_server.py
+++ b/update_server.py
@@ -1,34 +1,72 @@
 #!/usr/bin/env python3
-import hashlib
 from bottle import route, run, request, HTTPResponse, response
+import hashlib
+import json
 import pdb
+import os
 
 updated = True
 filename = "firmware.bin"
 
 @route('/update')
-def update():
+def updateRequest():
     #for header in request.headers.keys():
     #    print("{}: {}".format(header, request.headers.get(header, "")))
 
-    requestor_mac = request.headers.get("X-Esp8266-Sta-Mac")
-    running_fw_md5 = request.headers.get("X-Esp8266-Sketch-Md5", "")
-    print("Update request from {} with fw {}".format(requestor_mac, running_fw_md5))
-    available_fw_md5 = md5(filename)
+    if not "X-Esp8266-Sta-Mac" in request.headers:
+        return "Hello fellow friend! Looking for updates?"
 
-    print("Lates firmware available is {} - ".format(available_fw_md5), end="")
-    if available_fw_md5 != running_fw_md5:
-        print("Differs, updating…")
-        response.set_header('Content-Type', 'application/octet-stream')
-        response.set_header('Content-Disposition', ' attachment; filename={}'.format(filename))
-        with open(filename, "rb") as firmware:
-            fw = firmware.read()
-            response.set_header('Content-Length', '{}'.format(len(fw)))
-        response.set_header('X-MD5', available_fw_md5)
-        return fw
-    else:
+    readConfig()
+    global config
+    requestor_mac = request.headers.get("X-Esp8266-Sta-Mac", "")
+    running_fw_md5 = request.headers.get("X-Esp8266-Sketch-Md5", "")
+    role, version, fw_filename, fw_hash = getProperties(config, requestor_mac)
+    if role == "" and version == "" and fw_filename == "" and fw_hash == "":
+        print("No configuration for {} found. Serving 304 - No Firmware".format(requestor_mac))
+        return nofirmware()
+
+    print("Update request from {} with role {} (running fw hash: {}).".format(requestor_mac, role, running_fw_md5))
+    fw_path = os.path.join("firmwares", role, version, fw_filename)
+    try:
+        latest_fw_md5 = md5(fw_path)
+    except FileNotFoundError:
+        print("Configured firmware not found at {}".format(fw_path))
+        return nofirmware()
+
+    if latest_fw_md5 != fw_hash:
+        print("Configured hash ({}) does not match calculeted one ({})".format(fw_hash, latest_fw_md5))
+        return nofirmware()
+
+    print("Latest firmware for role \"{}\" is \"{}\" ({}) - ".format(role, fw_filename, latest_fw_md5), end="")
+    if latest_fw_md5 == running_fw_md5:
         print("Same, skipping…")
-        return HTTPResponse(status=304, body="")
+        return nofirmware()
+
+    print("Differs, updating…")
+    response.set_header('Content-Type', 'application/octet-stream')
+    response.set_header('Content-Disposition', ' attachment; filename={}'.format(filename))
+    with open(fw_path, "rb") as fw:
+        firmware = fw.read()
+        response.set_header('Content-Length', '{}'.format(len(firmware)))
+    response.set_header('X-MD5', latest_fw_md5)
+    print("Serving {} bytes to {} with role {}".format(len(firmware), requestor_mac, role))
+    return firmware
+
+def getProperties(config, mac):
+    configured_devices = config.get("devices", {}).keys()
+    devices = config.get("devices", {})
+    firmwares = config.get("firmwares", {})
+    match = devices.get(mac, {"role":"", "version":""})
+    role = match.get("role", "")
+    version = match.get("version", "")
+    role_firmwares = firmwares.get(role, {})
+    firmware = role_firmwares.get(version, {})
+    firmware_filename = firmware.get("filename", "")
+    firmware_hash = firmware.get("hash", "")
+    return role, version, firmware_filename, firmware_hash
+
+def nofirmware():
+    return HTTPResponse(status=304, body="")
 
 def md5(fname):
     hash_md5 = hashlib.md5()
@@ -37,4 +75,11 @@ def md5(fname):
             hash_md5.update(chunk)
     return hash_md5.hexdigest()
 
+def readConfig(config_name="config.json"):
+    global config
+    config = ""
+    with open(config_name, "r") as conf:
+        config = json.loads(conf.read())
+
+readConfig()
 run(host='localhost', port=8080, debug=True)