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

goosebit/updater/manager.py

@@ -1,82 +1,126 @@
 from __future__ import annotations
 
 import asyncio
+import re
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
-from typing import Callable
+from datetime import datetime
+from enum import StrEnum
+from typing import Callable, Optional
 
-import aiofiles
+from aiocache import Cache, cached
+from aiocache.serializers import PickleSerializer
 
-from goosebit.models import Device
-from goosebit.settings import POLL_TIME, SWUPDATE_FILES_DIR, TOKEN_SWU_DIR
-from goosebit.updater.misc import get_newest_fw
-from goosebit.updater.updates import FirmwareArtifact
+from goosebit.models import (
+    Device,
+    Firmware,
+    Hardware,
+    Rollout,
+    UpdateModeEnum,
+    UpdateStateEnum,
+)
+from goosebit.settings import POLL_TIME, POLL_TIME_UPDATING
+
+
+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.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_device_state(self, state: str) -> None:
+    async def update_hardware(self, hardware: Hardware) -> None:
+        return
+
+    async def update_device_state(self, state: UpdateStateEnum) -> None:
+        return
+
+    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_web_pwd(self, web_pwd: str) -> None:
+    async def update_name(self, name: str):
         return
 
-    async def update_last_ip(self, last_ip: str) -> None:
+    async def update_feed(self, feed: str):
         return
 
-    async def update_cfd_status(self, status: bool) -> None:
+    async def update_flavor(self, flavor: str):
         return
 
-    async def create_cfd_token(self) -> str:
-        return ""
+    async def update_config_data(self, **kwargs):
+        return
 
-    async def delete_cfd_token(self):
+    async def update_log_complete(self, log_complete: bool):
         return
 
+    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
 
-    async def publish_log(self, log_data: str | None):
-        for cb in self.log_subscribers:
-            await cb(log_data)
+    @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 cfd_provisioned(self) -> bool:
-        return False
+    def log_subscribers(self):
+        return UpdateManager.device_log_subscriptions.get(self.dev_id, [])
 
-    @abstractmethod
-    async def get_update_file(self) -> FirmwareArtifact: ...
+    @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_mode(self) -> str: ...
+    async def get_update(self) -> tuple[HandlingType, Firmware]: ...
 
     @abstractmethod
     async def update_log(self, log_data: str) -> None: ...
@@ -85,122 +129,221 @@ class UpdateManager(ABC):
 class UnknownUpdateManager(UpdateManager):
     def __init__(self, dev_id: str):
         super().__init__(dev_id)
-        self.poll_time = "00:00:05"
+        self.poll_time = POLL_TIME_UPDATING
 
-    async def get_update_file(self) -> FirmwareArtifact:
-        return FirmwareArtifact(get_newest_fw())
+    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 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.force_update = force_update
+        await self.save_device(device, update_fields=["force_update"])
 
     async def update_fw_version(self, version: str) -> None:
         device = await self.get_device()
         device.fw_version = version
+        await self.save_device(device, update_fields=["fw_version"])
 
-    async def update_device_state(self, state: str) -> None:
+    async def update_hardware(self, hardware: Hardware) -> None:
         device = await self.get_device()
-        device.last_state = state
+        device.hardware = hardware
+        await self.save_device(device, update_fields=["hardware"])
 
-    async def update_last_seen(self, last_seen: int) -> None:
+    async def update_device_state(self, state: UpdateStateEnum) -> None:
         device = await self.get_device()
-        device.last_seen = last_seen
+        device.last_state = state
+        await self.save_device(device, update_fields=["last_state"])
 
-    async def update_last_ip(self, last_ip: str) -> None:
+    async def update_last_connection(self, last_seen: int, last_ip: str) -> None:
         device = await self.get_device()
+        device.last_seen = last_seen
         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 get_update_file(self) -> FirmwareArtifact:
+    async def update_update(self, update_mode: UpdateModeEnum, firmware: Firmware | None):
         device = await self.get_device()
-        file = FirmwareArtifact(device.fw_file)
+        device.assigned_firmware = firmware
+        device.update_mode = update_mode
+        await self.save_device(device, update_fields=["assigned_firmware_id", "update_mode"])
 
-        if self.force_update:
-            return file
-        return file
+    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 get_update_mode(self) -> str:
+    async def update_feed(self, feed: str):
         device = await self.get_device()
+        device.feed = feed
+        await self.save_device(device, update_fields=["feed"])
 
-        file = await self.get_update_file()
-        if file.is_empty():
-            mode = "skip"
+    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.update_mode == UpdateModeEnum.ROLLOUT:
+            return (
+                await Rollout.filter(
+                    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_firmware(self) -> Firmware | None:
+        device = await self.get_device()
+
+        if device.update_mode == UpdateModeEnum.ROLLOUT:
+            rollout = await self.get_rollout()
+            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(self) -> tuple[HandlingType, Firmware]:
+        device = await self.get_device()
+        firmware = await self._get_firmware()
+
+        if firmware is None:
+            handling_type = HandlingType.SKIP
             self.poll_time = POLL_TIME
-        elif file.name == device.fw_version:
-            mode = "skip"
+
+        elif firmware.version == device.fw_version and not device.force_update:
+            handling_type = HandlingType.SKIP
             self.poll_time = POLL_TIME
-        else:
-            mode = "forced"
-            self.poll_time = "00:00:05"
 
-        if self.force_update:
-            mode = "forced"
-            self.poll_time = "00:00:05"
+        elif device.last_state == UpdateStateEnum.ERROR and not device.force_update:
+            handling_type = HandlingType.SKIP
+            self.poll_time = POLL_TIME
+
+        else:
+            handling_type = HandlingType.FORCED
+            self.poll_time = POLL_TIME_UPDATING
 
-        if mode == "forced" and self.update_complete:
-            self.update_complete = False
-            await self.clear_log()
+            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 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 self.save_device(
+            device,
+            update_fields=["progress", "last_log"],
+        )
+
     async def clear_log(self) -> None:
         device = await self.get_device()
         device.last_log = ""
+        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

@@ -3,22 +3,22 @@ import time
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.requests import Request
 
-from . import controller, download
-from .manager import get_update_manager_sync
+from goosebit.settings import TENANT
+
+from . import controller
+from .manager import get_update_manager
 
 
 async def verify_tenant(tenant: str):
-    if not tenant == "loadsync":
+    if not tenant == TENANT:
         raise HTTPException(404)
     return tenant
 
 
 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(
@@ -27,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()