ota.py

#!/usr/bin/python3

# TODO: read from header?
BROKER = "mqtt.c3pb.hack"
DEVICE = "my-device"


import sys
import paho.mqtt.client as mqtt


def on_connect(client, *ignored):
  client.subscribe("device/" + DEVICE + "/ota/progress")
  client.subscribe("device/" + DEVICE + "/ota/error")

  with open("build/qthing.bin", mode="rb") as f:
    client.publish("device/" + DEVICE + "/ota/$firmware", payload=f.read())

def on_message(client, userdata, msg):
  if msg.topic.endswith("error"):
    print("OTA failed: " + str(msg.payload))
    exit(1)

  progress = int(msg.payload)

  if progress == 0:
    print("OTA progress:")
  
  sys.stdout.write("\r[" + "#"*progress + " "*(100-progress) + "] " + str(progress) + "%")
  sys.stdout.flush()

  if progress == 100:
    print("\nOTA successful")
    if "--restart" in sys.argv:
      print("restarting device")
      client.publish("device/" + DEVICE + "/command", "restart")
    else:
      exit(0)
  
def on_publish(client, userdata, _ignored):
  if userdata.first_pub_done:
    exit(0)
  userdata.first_pub_done = True

class Userdata(object):
  first_pub_done = False

client = mqtt.Client(userdata=Userdata())
client.on_connect = on_connect
client.on_message = on_message
client.on_publish = on_publish

client.connect(BROKER)
client.loop_forever()