import requests, hashlib, os, time, yaml, subprocess, pathlib
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding

class OTAUpdater:
    def __init__(self, cfg):
        self.url = cfg["image_url"]
        self.pub_key = serialization.load_pem_public_key(open(cfg["public_key"],"rb").read())
        self.poll_s = cfg["poll_hours"] * 3600
        self.next = time.time() + self.poll_s
    def poll_if_due(self):
        if time.time() < self.next: return
        self.next += self.poll_s
        img = requests.get(self.url, timeout=10).content
        sha = hashlib.sha256(img).digest()
        sig  = requests.get(self.url + ".sig").content
        self.pub_key.verify(sig, sha, padding.PKCS1v15(), hashes.SHA256())
        pathlib.Path("/opt/rootfs.next.tar.gz").write_bytes(img)
        subprocess.run(["/usr/local/bin/swap_rootfs.sh"], check=True)