goosebit 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

goosebit/updater/manager.py

@@ -5,87 +5,122 @@ import re
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
 from datetime import datetime
+from enum import StrEnum
 from typing import Callable, Optional
 
-from goosebit.models import Device, Rollout
+from aiocache import Cache, cached
+from aiocache.serializers import PickleSerializer
+
+from goosebit.models import (
+    Device,
+    Firmware,
+    Hardware,
+    Rollout,
+    UpdateModeEnum,
+    UpdateStateEnum,
+)
 from goosebit.settings import POLL_TIME, POLL_TIME_UPDATING
-from goosebit.updates.artifacts import FirmwareArtifact
+
+
+class HandlingType(StrEnum):
+    SKIP = "skip"
+    ATTEMPT = "attempt"
+    FORCED = "forced"
 
 
 class UpdateManager(ABC):
+    device_log_subscriptions: dict[str, list[Callable]] = {}
+    device_poll_time: dict[str, str] = {}
+
     def __init__(self, dev_id: str):
         self.dev_id = dev_id
-        self.config_data = {}
-        self.device = None
-        self.force_update = False
-        self.update_complete = False
-        self.poll_time = POLL_TIME
-        self.log_subscribers: list[Callable] = []
 
     async def get_device(self) -> Device | None:
         return
 
-    async def save(self) -> None:
+    async def update_force_update(self, force_update: bool) -> None:
         return
 
     async def update_fw_version(self, version: str) -> None:
         return
 
-    async def update_hw_model(self, hw_model: str) -> None:
+    async def update_hardware(self, hardware: Hardware) -> None:
         return
 
-    async def update_hw_revision(self, hw_revision: str) -> None:
+    async def update_device_state(self, state: UpdateStateEnum) -> None:
         return
 
-    async def update_device_state(self, state: str) -> None:
+    async def update_last_connection(self, last_seen: int, last_ip: str) -> None:
         return
 
-    async def update_last_seen(self, last_seen: int) -> None:
+    async def update_update(self, update_mode: UpdateModeEnum, firmware: Firmware | None):
         return
 
-    async def update_last_ip(self, last_ip: str) -> None:
+    async def update_name(self, name: str):
         return
 
-    async def get_rollout(self) -> Optional[Rollout]:
-        return None
+    async def update_feed(self, feed: str):
+        return
+
+    async def update_flavor(self, flavor: str):
+        return
 
     async def update_config_data(self, **kwargs):
-        await self.update_hw_model(kwargs.get("hw_model") or "default")
-        await self.update_hw_revision(kwargs.get("hw_revision") or "default")
+        return
 
-        device = await self.get_device()
-        if device.last_state == "unknown":
-            await self.update_device_state("registered")
-        await self.save()
+    async def update_log_complete(self, log_complete: bool):
+        return
 
-        self.config_data.update(kwargs)
+    async def get_rollout(self) -> Optional[Rollout]:
+        return None
 
     @asynccontextmanager
     async def subscribe_log(self, callback: Callable):
         device = await self.get_device()
-        self.log_subscribers.append(callback)
+        subscribers = self.log_subscribers
+        subscribers.append(callback)
+        self.log_subscribers = subscribers
         await callback(device.last_log)
         try:
             yield
         except asyncio.CancelledError:
             pass
         finally:
-            self.log_subscribers.remove(callback)
+            subscribers = self.log_subscribers
+            subscribers.remove(callback)
+            self.log_subscribers = subscribers
 
     @property
     def poll_seconds(self):
         time_obj = datetime.strptime(self.poll_time, "%H:%M:%S")
         return time_obj.hour * 3600 + time_obj.minute * 60 + time_obj.second
 
+    @property
+    def log_subscribers(self):
+        return UpdateManager.device_log_subscriptions.get(self.dev_id, [])
+
+    @log_subscribers.setter
+    def log_subscribers(self, value: list):
+        UpdateManager.device_log_subscriptions[self.dev_id] = value
+
+    @property
+    def poll_time(self):
+        return UpdateManager.device_poll_time.get(self.dev_id, POLL_TIME)
+
+    @poll_time.setter
+    def poll_time(self, value: str):
+        if not value == POLL_TIME:
+            UpdateManager.device_poll_time[self.dev_id] = value
+            return
+        if self.dev_id in UpdateManager.device_poll_time:
+            del UpdateManager.device_poll_time[self.dev_id]
+
     async def publish_log(self, log_data: str | None):
         for cb in self.log_subscribers:
             await cb(log_data)
 
     @abstractmethod
-    async def get_update_file(self) -> FirmwareArtifact: ...
-
-    @abstractmethod
-    async def get_update_mode(self) -> str: ...
+    async def get_update(self) -> tuple[HandlingType, Firmware]: ...
 
     @abstractmethod
     async def update_log(self, log_data: str) -> None: ...
@@ -96,152 +131,219 @@ class UnknownUpdateManager(UpdateManager):
         super().__init__(dev_id)
         self.poll_time = POLL_TIME_UPDATING
 
-    async def get_update_file(self) -> FirmwareArtifact:
-        return FirmwareArtifact("latest")
+    async def _get_firmware(self) -> Firmware:
+        return await Firmware.latest(await self.get_device())
 
-    async def get_update_mode(self) -> str:
-        return "forced"
+    async def get_update(self) -> tuple[HandlingType, Firmware]:
+        firmware = await self._get_firmware()
+        return HandlingType.FORCED, firmware
 
     async def update_log(self, log_data: str) -> None:
         return
 
 
 class DeviceUpdateManager(UpdateManager):
+    hardware_default = None
+
+    @cached(
+        ttl=600,
+        key_builder=lambda fn, self: self.dev_id,
+        cache=Cache.MEMORY,
+        serializer=PickleSerializer(),
+        namespace="main",
+    )
     async def get_device(self) -> Device:
-        if self.device:
-            return self.device
-        self.device = (await Device.get_or_create(uuid=self.dev_id))[0]
-        return self.device
+        hardware = DeviceUpdateManager.hardware_default
+        if hardware is None:
+            hardware = (await Hardware.get_or_create(model="default", revision="default"))[0]
+            DeviceUpdateManager.hardware_default = hardware
 
-    async def save(self) -> None:
-        await self.device.save()
+        return (await Device.get_or_create(uuid=self.dev_id, defaults={"hardware": hardware}))[0]
 
-    async def update_fw_version(self, version: str) -> None:
+    async def save_device(self, device: Device, update_fields: list[str]):
+        cache = Cache(namespace="main")
+        await cache.set(self.dev_id, device, ttl=600)
+        await device.save(update_fields=update_fields)
+
+    async def update_force_update(self, force_update: bool) -> None:
         device = await self.get_device()
-        device.fw_version = version
+        device.force_update = force_update
+        await self.save_device(device, update_fields=["force_update"])
 
-    async def update_hw_model(self, hw_model: str) -> None:
+    async def update_fw_version(self, version: str) -> None:
         device = await self.get_device()
-        device.hw_model = hw_model
+        device.fw_version = version
+        await self.save_device(device, update_fields=["fw_version"])
 
-    async def update_hw_revision(self, hw_revision: str) -> None:
+    async def update_hardware(self, hardware: Hardware) -> None:
         device = await self.get_device()
-        device.hw_revision = hw_revision
+        device.hardware = hardware
+        await self.save_device(device, update_fields=["hardware"])
 
-    async def update_device_state(self, state: str) -> None:
+    async def update_device_state(self, state: UpdateStateEnum) -> None:
         device = await self.get_device()
         device.last_state = state
+        await self.save_device(device, update_fields=["last_state"])
 
-    async def update_last_seen(self, last_seen: int) -> None:
+    async def update_last_connection(self, last_seen: int, last_ip: str) -> None:
         device = await self.get_device()
         device.last_seen = last_seen
-
-    async def update_last_ip(self, last_ip: str) -> None:
-        device = await self.get_device()
         if ":" in last_ip:
             device.last_ipv6 = last_ip
+            await self.save_device(device, update_fields=["last_seen", "last_ipv6"])
         else:
             device.last_ip = last_ip
+            await self.save_device(device, update_fields=["last_seen", "last_ip"])
+
+    async def update_update(self, update_mode: UpdateModeEnum, firmware: Firmware | None):
+        device = await self.get_device()
+        device.assigned_firmware = firmware
+        device.update_mode = update_mode
+        await self.save_device(device, update_fields=["assigned_firmware_id", "update_mode"])
+
+    async def update_name(self, name: str):
+        device = await self.get_device()
+        device.name = name
+        await self.save_device(device, update_fields=["name"])
+
+    async def update_feed(self, feed: str):
+        device = await self.get_device()
+        device.feed = feed
+        await self.save_device(device, update_fields=["feed"])
+
+    async def update_flavor(self, flavor: str):
+        device = await self.get_device()
+        device.flavor = flavor
+        await self.save_device(device, update_fields=["flavor"])
+
+    async def update_config_data(self, **kwargs):
+        model = kwargs.get("hw_model") or "default"
+        revision = kwargs.get("hw_revision") or "default"
+        hardware = (await Hardware.get_or_create(model=model, revision=revision))[0]
+        device = await self.get_device()
+        modified = False
+
+        if device.hardware != hardware:
+            device.hardware = hardware
+            modified = True
+
+        if device.last_state == UpdateStateEnum.UNKNOWN:
+            device.last_state = UpdateStateEnum.REGISTERED
+            modified = True
+
+        if modified:
+            await self.save_device(device, update_fields=["hardware_id", "last_state"])
+
+    async def update_log_complete(self, log_complete: bool):
+        device = await self.get_device()
+        device.log_complete = log_complete
+        await self.save_device(device, update_fields=["log_complete"])
 
     async def get_rollout(self) -> Optional[Rollout]:
         device = await self.get_device()
 
-        if device.fw_file == "none":
+        if device.update_mode == UpdateModeEnum.ROLLOUT:
             return (
                 await Rollout.filter(
-                    hw_model=device.hw_model,
-                    hw_revision=device.hw_revision,
                     feed=device.feed,
                     flavor=device.flavor,
+                    paused=False,
+                    firmware__compatibility__devices__uuid=device.uuid,
                 )
                 .order_by("-created_at")
                 .first()
+                .prefetch_related("firmware")
             )
 
         return None
 
-    async def get_update_file(self) -> FirmwareArtifact:
+    async def _get_firmware(self) -> Firmware | None:
         device = await self.get_device()
-        file = device.fw_file
 
-        if file == "none":
+        if device.update_mode == UpdateModeEnum.ROLLOUT:
             rollout = await self.get_rollout()
-            if rollout and not rollout.paused:
-                file = rollout.fw_file
-
-        return FirmwareArtifact(file, device.hw_model, device.hw_revision)
+            if not rollout or rollout.paused:
+                return None
+            await rollout.fetch_related("firmware")
+            return rollout.firmware
+        if device.update_mode == UpdateModeEnum.ASSIGNED:
+            await device.fetch_related("assigned_firmware")
+            return device.assigned_firmware
+
+        if device.update_mode == UpdateModeEnum.LATEST:
+            return await Firmware.latest(device)
+
+        assert device.update_mode == UpdateModeEnum.PINNED
+        return None
 
-    async def get_update_mode(self) -> str:
+    async def get_update(self) -> tuple[HandlingType, Firmware]:
         device = await self.get_device()
+        firmware = await self._get_firmware()
 
-        file = await self.get_update_file()
-        if file.is_empty():
-            mode = "skip"
+        if firmware is None:
+            handling_type = HandlingType.SKIP
             self.poll_time = POLL_TIME
-        elif file.name == device.fw_version and not self.force_update:
-            mode = "skip"
+
+        elif firmware.version == device.fw_version and not device.force_update:
+            handling_type = HandlingType.SKIP
             self.poll_time = POLL_TIME
-        elif device.last_state == "error" and not self.force_update:
-            mode = "skip"
+
+        elif device.last_state == UpdateStateEnum.ERROR and not device.force_update:
+            handling_type = HandlingType.SKIP
             self.poll_time = POLL_TIME
+
         else:
-            mode = "forced"
+            handling_type = HandlingType.FORCED
             self.poll_time = POLL_TIME_UPDATING
 
-            if self.update_complete:
-                self.update_complete = False
+            if device.log_complete:
+                await self.update_log_complete(False)
                 await self.clear_log()
 
-        return mode
+        return handling_type, firmware
 
     async def update_log(self, log_data: str) -> None:
         if log_data is None:
             return
         device = await self.get_device()
+
+        if device.last_log is None:
+            device.last_log = ""
+
         matches = re.findall(r"Downloaded (\d+)%", log_data)
         if matches:
             device.progress = matches[-1]
-        if device.last_log is None:
-            device.last_log = ""
+
         if log_data.startswith("Installing Update Chunk Artifacts."):
-            await self.clear_log()
-        if log_data == "All Chunks Installed.":
-            self.force_update = False
-            self.update_complete = True
+            # clear log
+            device.last_log = ""
+            await self.publish_log(None)
+
         if not log_data == "Skipped Update.":
             device.last_log += f"{log_data}\n"
             await self.publish_log(f"{log_data}\n")
-        await device.save()
+
+        await self.save_device(
+            device,
+            update_fields=["progress", "last_log"],
+        )
 
     async def clear_log(self) -> None:
         device = await self.get_device()
         device.last_log = ""
-        await device.save()
+        await self.save_device(device, update_fields=["last_log"])
         await self.publish_log(None)
 
 
-device_managers = {"unknown": UnknownUpdateManager("unknown")}
+async def get_update_manager(dev_id: str) -> UpdateManager:
+    if dev_id == "unknown":
+        return UnknownUpdateManager("unknown")
+    else:
+        return DeviceUpdateManager(dev_id)
 
 
-async def get_update_manager(dev_id: str) -> UpdateManager:
-    global device_managers
-    if device_managers.get(dev_id) is None:
-        device_managers[dev_id] = DeviceUpdateManager(dev_id)
-    return device_managers[dev_id]
-
-
-def get_update_manager_sync(dev_id: str) -> UpdateManager:
-    global device_managers
-    if device_managers.get(dev_id) is None:
-        device_managers[dev_id] = DeviceUpdateManager(dev_id)
-    return device_managers[dev_id]
-
-
-async def delete_device(dev_id: str) -> None:
-    global device_managers
-    try:
-        updater = get_update_manager_sync(dev_id)
-        await (await updater.get_device()).delete()
-        del device_managers[dev_id]
-    except KeyError:
-        pass
+async def delete_devices(ids: list[str]):
+    await Device.filter(id__in=ids).delete()
+    cache = Cache(namespace="main")
+    for dev_id in ids:
+        await cache.delete(dev_id)

goosebit/updater/routes.py

@@ -5,8 +5,8 @@ from fastapi.requests import Request
 
 from goosebit.settings import TENANT
 
-from . import controller, download
-from .manager import get_update_manager_sync
+from . import controller
+from .manager import get_update_manager
 
 
 async def verify_tenant(tenant: str):
@@ -17,10 +17,8 @@ async def verify_tenant(tenant: str):
 
 async def log_last_connection(request: Request, dev_id: str):
     host = request.client.host
-    updater = get_update_manager_sync(dev_id)
-    await updater.update_last_ip(host)
-    await updater.update_last_seen(round(time.time()))
-    await updater.save()
+    updater = await get_update_manager(dev_id)
+    await updater.update_last_connection(round(time.time()), host)
 
 
 router = APIRouter(
@@ -29,4 +27,3 @@ router = APIRouter(
     tags=["ddi"],
 )
 router.include_router(controller.router)
-router.include_router(download.router)

goosebit/updates/__init__.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+from urllib.parse import unquote, urlparse
+from urllib.request import url2pathname
+
+from fastapi.requests import Request
+
+from goosebit.models import Firmware, Hardware
+
+from . import swdesc
+
+
+async def create_firmware_update(uri: str):
+    parsed_uri = urlparse(uri)
+
+    if parsed_uri.scheme == "file":
+        file_path = Path(url2pathname(unquote(parsed_uri.path)))
+        update_info = await swdesc.parse_file(file_path)
+        if update_info is None:
+            return
+    elif parsed_uri.scheme.startswith("http"):
+        update_info = await swdesc.parse_remote(uri)
+        if update_info is None:
+            return
+    else:
+        return
+
+    firmware = (
+        await Firmware.get_or_create(
+            uri=uri,
+            version=str(update_info["version"]),
+            size=update_info["size"],
+            hash=update_info["hash"],
+        )
+    )[0]
+
+    for comp in update_info["compatibility"]:
+        model = comp.get("hw_model", "default")
+        revision = comp.get("hw_revision", "default")
+        await firmware.compatibility.add((await Hardware.get_or_create(model=model, revision=revision))[0])
+    await firmware.save()
+    return firmware
+
+
+def generate_chunk(request: Request, firmware: Firmware | None) -> list:
+    if firmware is None:
+        return []
+    if firmware.local:
+        href = str(
+            request.url_for(
+                "download_file",
+                file_id=firmware.id,
+            )
+        )
+    else:
+        href = firmware.uri
+    return [
+        {
+            "part": "os",
+            "version": "1",
+            "name": firmware.path.name,
+            "artifacts": [
+                {
+                    "filename": firmware.path.name,
+                    "hashes": {"sha1": firmware.hash},
+                    "size": firmware.size,
+                    "_links": {"download": {"href": href}},
+                }
+            ],
+        }
+    ]

goosebit/updates/swdesc.py

@@ -0,0 +1,83 @@
+import hashlib
+import logging
+from pathlib import Path
+from typing import Any
+
+import aiofiles
+import httpx
+import libconf
+import semver
+
+from goosebit.settings import UPDATES_DIR
+
+logger = logging.getLogger(__name__)
+
+
+def _append_compatibility(boardname, value, compatibility):
+    if "hardware-compatibility" in value:
+        for revision in value["hardware-compatibility"]:
+            compatibility.append({"hw_model": boardname, "hw_revision": revision})
+
+
+def parse_descriptor(swdesc: libconf.AttrDict[Any, Any | None]):
+    swdesc_attrs = {}
+    try:
+        swdesc_attrs["version"] = semver.Version.parse(swdesc["software"]["version"])
+        compatibility = []
+        _append_compatibility("default", swdesc["software"], compatibility)
+
+        for key in swdesc["software"]:
+            element = swdesc["software"][key]
+            _append_compatibility(key, element, compatibility)
+
+            if isinstance(element, dict):
+                for key2 in element:
+                    _append_compatibility(key, element[key2], compatibility)
+
+        swdesc_attrs["compatibility"] = compatibility
+    except KeyError as e:
+        logging.warning(f"Parsing swu descriptor failed, error={e}")
+        raise ValueError("Parsing swu descriptor failed", e)
+
+    return swdesc_attrs
+
+
+async def parse_file(file: Path):
+    async with aiofiles.open(file, "r+b") as f:
+        # get file size
+        size = int((await f.read(110))[54:62], 16)
+        filename = b""
+        next_byte = await f.read(1)
+        while not next_byte == b"\x00":
+            filename += next_byte
+            next_byte = await f.read(1)
+        # 4 null bytes
+        await f.read(3)
+
+        # should always be the first file
+        if not filename == b"sw-description":
+            return None
+
+        swdesc = libconf.loads((await f.read(size)).decode("utf-8"))
+
+        swdesc_attrs = parse_descriptor(swdesc)
+        swdesc_attrs["size"] = file.stat().st_size
+        swdesc_attrs["hash"] = _sha1_hash_file(file)
+        return swdesc_attrs
+
+
+async def parse_remote(url: str):
+    async with httpx.AsyncClient() as c:
+        file = await c.get(url)
+        temp_file = UPDATES_DIR.joinpath("temp.swu")
+        async with aiofiles.open(temp_file, "w+b") as f:
+            await f.write(file.content)
+    parsed_file = await parse_file(temp_file)
+    temp_file.unlink()
+    return parsed_file
+
+
+def _sha1_hash_file(file_path: Path):
+    with file_path.open("rb") as f:
+        sha1_hash = hashlib.file_digest(f, "sha1")
+    return sha1_hash.hexdigest()