goosebit 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

goosebit/device_manager.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import asyncio
+import re
+from enum import StrEnum
+from typing import Any, Awaitable, Callable, Optional
+
+from aiocache import caches
+from fastapi.requests import Request
+
+from goosebit.db.models import (
+    Device,
+    Hardware,
+    Rollout,
+    Software,
+    UpdateModeEnum,
+    UpdateStateEnum,
+)
+from goosebit.schema.updates import UpdateChunk
+
+caches.set_config(
+    {
+        "default": {
+            "cache": "aiocache.SimpleMemoryCache",
+            "serializer": {"class": "aiocache.serializers.PickleSerializer"},
+            "ttl": 600,
+        },
+    }
+)
+
+
+class HandlingType(StrEnum):
+    SKIP = "skip"
+    ATTEMPT = "attempt"
+    FORCED = "forced"
+
+
+class DeviceManager:
+    _hardware_default = None
+
+    _update_sources: list[Callable[[Request, Device], Awaitable[tuple[HandlingType, UpdateChunk | None]]]] = []
+    _config_callbacks: list[Callable[[Device, dict[str, Any]], Awaitable[None]]] = []
+
+    @staticmethod
+    async def get_device(dev_id: str) -> Device:
+        cache = caches.get("default")
+        device = await cache.get(dev_id)
+        if device:
+            return device
+
+        hardware = DeviceManager._hardware_default
+        if hardware is None:
+            hardware = (await Hardware.get_or_create(model="default", revision="default"))[0]
+            DeviceManager._hardware_default = hardware
+
+        device = (await Device.get_or_create(id=dev_id, defaults={"hardware": hardware}))[0]
+        result = await cache.set(device.id, device, ttl=600)
+        assert result, "device being cached"
+
+        return device
+
+    @staticmethod
+    async def save_device(device: Device, update_fields: list[str]):
+        await device.save(update_fields=update_fields)
+
+        # only update cache after a successful database save
+        result = await caches.get("default").set(device.id, device, ttl=600)
+        assert result, "device being cached"
+
+    @staticmethod
+    async def update_auth_token(device: Device, auth_token: str) -> None:
+        device.auth_token = auth_token
+        await DeviceManager.save_device(device, update_fields=["auth_token"])
+
+    @staticmethod
+    async def update_force_update(device: Device, force_update: bool) -> None:
+        device.force_update = force_update
+        await DeviceManager.save_device(device, update_fields=["force_update"])
+
+    @staticmethod
+    async def update_sw_version(device: Device, version: str) -> None:
+        device.sw_version = version
+        await DeviceManager.save_device(device, update_fields=["sw_version"])
+
+    @staticmethod
+    async def update_hardware(device: Device, hardware: Hardware) -> None:
+        device.hardware = hardware
+        await DeviceManager.save_device(device, update_fields=["hardware"])
+
+    @staticmethod
+    async def update_device_state(device: Device, state: UpdateStateEnum) -> None:
+        device.last_state = state
+        await DeviceManager.save_device(device, update_fields=["last_state"])
+
+    @staticmethod
+    async def update_last_connection(device: Device, last_seen: int, last_ip: str | None = None) -> None:
+        device.last_seen = last_seen
+        if last_ip is None:
+            await DeviceManager.save_device(device, update_fields=["last_seen"])
+        elif ":" in last_ip:
+            device.last_ipv6 = last_ip
+            await DeviceManager.save_device(device, update_fields=["last_seen", "last_ipv6"])
+        else:
+            device.last_ip = last_ip
+            await DeviceManager.save_device(device, update_fields=["last_seen", "last_ip"])
+
+    @staticmethod
+    async def update_update(device: Device, update_mode: UpdateModeEnum, software: Software | None):
+        device.assigned_software = software
+        device.update_mode = update_mode
+        if not update_mode == UpdateModeEnum.ROLLOUT:
+            device.feed = None
+        await DeviceManager.save_device(device, update_fields=["assigned_software_id", "update_mode", "feed"])
+
+    @staticmethod
+    async def update_name(device: Device, name: str):
+        device.name = name
+        await DeviceManager.save_device(device, update_fields=["name"])
+
+    @staticmethod
+    async def update_feed(device: Device, feed: str):
+        device.feed = feed
+        await DeviceManager.save_device(device, update_fields=["feed"])
+
+    @staticmethod
+    def add_config_callback(callback: Callable[[Device, dict[str, Any]], Awaitable[None]]):
+        DeviceManager._config_callbacks.append(callback)
+
+    @staticmethod
+    def remove_config_callback(callback: Callable[[Device, dict[str, Any]], Awaitable[None]]):
+        DeviceManager._config_callbacks.remove(callback)
+
+    @staticmethod
+    async def update_config_data(device: Device, **kwargs: dict[str, Any]):
+        model = kwargs.get("hw_boardname") or "default"
+        revision = kwargs.get("hw_revision") or "default"
+        sw_version = kwargs.get("sw_version")
+
+        await asyncio.gather(
+            *[cb(device, **kwargs) for cb in DeviceManager._config_callbacks]  # type: ignore[call-arg]
+        )
+
+        hardware = (await Hardware.get_or_create(model=model, revision=revision))[0]
+        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 device.sw_version != sw_version:
+            device.sw_version = sw_version
+            modified = True
+
+        if modified:
+            await DeviceManager.save_device(device, update_fields=["hardware_id", "last_state", "sw_version"])
+
+    @staticmethod
+    async def deployment_action_start(device: Device):
+        device.last_log = ""
+        device.progress = 0
+        await DeviceManager.save_device(device, update_fields=["last_log", "progress"])
+
+    @staticmethod
+    async def deployment_action_success(device: Device):
+        device.progress = 100
+        await DeviceManager.save_device(device, update_fields=["progress"])
+
+    @staticmethod
+    async def get_rollout(device: Device) -> Rollout | None:
+        if device.update_mode == UpdateModeEnum.ROLLOUT:
+            return (
+                await Rollout.filter(
+                    feed=device.feed,
+                    paused=False,
+                    software__compatibility__devices__id=device.id,
+                )
+                .order_by("-created_at")
+                .first()
+                .prefetch_related("software")
+            )
+
+        return None
+
+    @staticmethod
+    async def _get_software(device: Device) -> Software | None:
+        if device.update_mode == UpdateModeEnum.ROLLOUT:
+            rollout = await DeviceManager.get_rollout(device)
+            if not rollout or rollout.paused:
+                return None
+            await rollout.fetch_related("software")
+            return rollout.software
+        if device.update_mode == UpdateModeEnum.ASSIGNED:
+            await device.fetch_related("assigned_software")
+            return device.assigned_software
+
+        if device.update_mode == UpdateModeEnum.LATEST:
+            return await Software.latest(device)
+
+        assert device.update_mode == UpdateModeEnum.PINNED
+        return None
+
+    @staticmethod
+    def add_update_source(source: Callable[[Request, Device], Awaitable[tuple[HandlingType, UpdateChunk | None]]]):
+        DeviceManager._update_sources.append(source)
+
+    @staticmethod
+    async def get_alt_src_updates(request: Request, device: Device) -> list[tuple[HandlingType, UpdateChunk | None]]:
+        return await asyncio.gather(*[source(request, device) for source in DeviceManager._update_sources])
+
+    @staticmethod
+    async def get_update(device: Device) -> tuple[HandlingType, Software | None]:
+        software = await DeviceManager._get_software(device)
+
+        if software is None:
+            handling_type = HandlingType.SKIP
+
+        elif software.version == device.sw_version and not device.force_update:
+            handling_type = HandlingType.SKIP
+
+        elif device.last_state == UpdateStateEnum.ERROR and not device.force_update:
+            handling_type = HandlingType.SKIP
+
+        else:
+            handling_type = HandlingType.FORCED
+
+        return handling_type, software
+
+    @staticmethod
+    async def update_log(device: Device, log_data: str) -> None:
+        if log_data is None:
+            return
+
+        if device.last_log is None:
+            device.last_log = ""
+
+        # SWUpdate-specific log parsing to report progress
+        matches = re.findall(r"Downloaded (\d+)%", log_data)
+        if matches:
+            device.progress = matches[-1]
+
+        device.last_log += f"{log_data}\n"
+
+        await DeviceManager.save_device(device, update_fields=["progress", "last_log"])
+
+    @staticmethod
+    async def delete_devices(ids: list[str]):
+        await Device.filter(id__in=ids).delete()
+        for dev_id in ids:
+            result = await caches.get("default").delete(dev_id)
+            assert result == 1, "device has been cached"
+
+
+async def get_device(dev_id: str) -> Device:
+    return await DeviceManager.get_device(dev_id)
+
+
+async def get_device_or_none(dev_id: str) -> Optional[Device]:
+    return await Device.get_or_none(id=dev_id)

goosebit/plugins/__init__.py

@@ -0,0 +1,32 @@
+import logging
+from importlib.metadata import entry_points
+
+from goosebit.schema.plugins import PluginSchema
+from goosebit.settings import config
+
+logger = logging.getLogger(__name__)
+
+
+def load() -> list:
+    plugin_configs = []
+    logger.info("Checking for plugins to be loaded...")
+    if len(config.plugins) == 0:
+        logger.info("No plugins found.")
+        return []
+    logger.info(f"Found plugins enabled in config: {config.plugins}")
+
+    entries = entry_points(group="goosebit.plugins")
+
+    for plugin in config.plugins:
+        modules = entries.select(name=plugin)
+        for module in modules:  # should be 1 or 0
+            logger.info(f"Loading plugin: {plugin}")
+            loaded_plugin = module.load()
+            if not hasattr(loaded_plugin, "config"):
+                logger.error(f"Failed to load plugin: {plugin}, plugin has not defined config")
+                continue
+            if not isinstance(loaded_plugin.config, PluginSchema):
+                logger.error(f"Failed to load plugin: {plugin}, config is not an instance of PluginSchema")
+                continue
+            plugin_configs.append(loaded_plugin.config)
+    return plugin_configs

goosebit/schema/devices.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import time
+from datetime import datetime
 from enum import Enum, IntEnum, StrEnum
 from typing import Annotated
 
@@ -8,7 +9,7 @@ from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, computed_fie
 
 from goosebit.db.models import UpdateModeEnum, UpdateStateEnum
 from goosebit.schema.software import HardwareSchema, SoftwareSchema
-from goosebit.updater.manager import DeviceUpdateManager
+from goosebit.settings import config
 
 
 class ConvertableEnum(StrEnum):
@@ -29,14 +30,14 @@ UpdateModeSchema = enum_factory("UpdateModeSchema", UpdateModeEnum)
 class DeviceSchema(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
-    uuid: str
+    id: str
     name: str | None
     sw_version: str | None
 
     assigned_software: SoftwareSchema | None = Field(exclude=True)
     hardware: HardwareSchema | None = Field(exclude=True)
 
-    feed: str
+    feed: str | None
     progress: int | None
     last_state: Annotated[UpdateStateSchema, BeforeValidator(UpdateStateSchema.convert)]  # type: ignore[valid-type]
     update_mode: Annotated[UpdateModeSchema, BeforeValidator(UpdateModeSchema.convert)]  # type: ignore[valid-type]
@@ -45,10 +46,11 @@ class DeviceSchema(BaseModel):
     last_seen: Annotated[
         int | None, BeforeValidator(lambda last_seen: round(time.time() - last_seen) if last_seen is not None else None)
     ]
+    auth_token: str | None
 
     @computed_field  # type: ignore[misc]
     @property
-    def online(self) -> bool | None:
+    def polling(self) -> bool | None:
         return self.last_seen < (self.poll_seconds + 10) if self.last_seen is not None else None
 
     @computed_field  # type: ignore[misc]
@@ -74,4 +76,5 @@ class DeviceSchema(BaseModel):
     @computed_field  # type: ignore[misc]
     @property
     def poll_seconds(self) -> int:
-        return DeviceUpdateManager(self.uuid).poll_seconds
+        time_obj = datetime.strptime(config.poll_time, "%H:%M:%S")
+        return time_obj.hour * 3600 + time_obj.minute * 60 + time_obj.second

goosebit/schema/plugins.py

@@ -0,0 +1,67 @@
+import inspect
+from typing import Any, Awaitable, Callable, Type
+
+from fastapi import APIRouter
+from fastapi.requests import Request
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from pydantic import BaseModel, Field, computed_field
+from pydantic_settings import (
+    BaseSettings,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+    YamlConfigSettingsSource,
+)
+
+from goosebit.db import Device
+from goosebit.device_manager import HandlingType
+from goosebit.schema.updates import UpdateChunk
+from goosebit.settings import config
+
+
+def get_module_name():
+    module = inspect.getmodule(inspect.stack()[2][0])
+    if module is not None:
+        return module.__name__.split(".")[0]
+    raise TypeError("Could not discover plugin module name")
+
+
+class PluginSchema(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    name: str = Field(
+        default_factory=get_module_name
+    )  # get the name of the package this was initialized in (plugin package)
+    middleware: Type[Any] | None = None  # ASGI middleware class
+    router: APIRouter | None = None
+    db_model_path: str | None = None
+    static_files: StaticFiles | None = None
+    templates: Jinja2Templates | None = None
+    update_source_hook: Callable[[Request, Device], Awaitable[tuple[HandlingType, UpdateChunk | None]]] | None = None
+    config_data_hook: Callable[[Device, dict[str, Any]], Awaitable[None]] | None = None
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def url_prefix(self) -> str:
+        return f"/plugins/{self.name}"
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def static_files_name(self) -> str:
+        return f"{self.name}_static"
+
+
+class PluginSettings(BaseSettings):
+    model_config = SettingsConfigDict(extra="ignore")
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return tuple([env_settings, YamlConfigSettingsSource(settings_cls, config.config_file)])

goosebit/schema/updates.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+
+
+class UpdateChunkArtifact(BaseModel):
+    filename: str
+    hashes: dict[str, str]
+    size: int
+    links: dict[str, dict[str, str]] = Field(serialization_alias="_links")
+
+
+class UpdateChunk(BaseModel):
+    part: str = "os"
+    version: str = "1"
+    name: str
+    artifacts: list[UpdateChunkArtifact]

goosebit/schema/users.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class UserSchema(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    username: str
+    enabled: bool
+    permissions: list[str]

goosebit/settings/__init__.py

@@ -12,6 +12,3 @@ logger = logging.getLogger(__name__)
 
 if config.config_file is not None:
     logger.info(f"Loading settings from: {config.config_file}")
-
-
-USERS = {u.username: u for u in config.users}

goosebit/settings/schema.py

@@ -1,10 +1,11 @@
 import os
 import secrets
+from enum import StrEnum
 from pathlib import Path
 from typing import Annotated
 
-from joserfc.rfc7518.oct_key import OctKey
-from pydantic import BaseModel, BeforeValidator, Field
+from joserfc.jwk import OctKey
+from pydantic import BaseModel, BeforeValidator, Field, model_validator
 from pydantic_settings import (
     BaseSettings,
     PydanticBaseSettingsSource,
@@ -12,16 +13,36 @@ from pydantic_settings import (
     YamlConfigSettingsSource,
 )
 
-from .const import CURRENT_DIR, GOOSEBIT_ROOT_DIR, LOGGING_DEFAULT, PWD_CXT
+from .const import CURRENT_DIR, GOOSEBIT_ROOT_DIR, LOGGING_DEFAULT
 
 
-class User(BaseModel):
-    username: str
-    hashed_pwd: Annotated[str, BeforeValidator(PWD_CXT.hash)] = Field(validation_alias="password")
-    permissions: set[str]
+class DeviceAuthMode(StrEnum):
+    SETUP = "setup"  # setup mode, any devices polling with an auth token that don't have one will save it
+    STRICT = "strict"  # all devices must have keys, and all keys must be set up with the API
+    EXTERNAL = "external"  # all devices must have keys, keys must be validated by external auth service
+    LAX = "lax"  # devices may or may not use keys, but device with keys set must poll with them
 
-    def get_json_permissions(self):
-        return [str(p) for p in self.permissions]
+
+class ExternalAuthMode(StrEnum):
+    BEARER = "bearer"  # validate token with external service using Bearer token
+    JSON = "json"  # validate token with external service using JSON payload
+
+
+class DeviceAuthSettings(BaseModel):
+    enable: bool = False
+    mode: DeviceAuthMode = DeviceAuthMode.STRICT
+    external_url: str | None = None
+    external_mode: ExternalAuthMode = ExternalAuthMode.JSON
+    external_json_key: str = "token"
+
+    @model_validator(mode="after")
+    def validate_external_mode_config(self):
+        if self.mode == DeviceAuthMode.EXTERNAL:
+            if self.external_url is None:
+                raise ValueError("External URL is required when using external authentication mode")
+            if self.external_mode == ExternalAuthMode.JSON and not self.external_json_key:
+                raise ValueError("External JSON key is required when using JSON mode")
+        return self
 
 
 class PrometheusSettings(BaseModel):
@@ -32,22 +53,47 @@ class MetricsSettings(BaseModel):
     prometheus: PrometheusSettings = PrometheusSettings()
 
 
+class StorageType(StrEnum):
+    FILESYSTEM = "filesystem"
+    S3 = "s3"
+
+
+class S3StorageSettings(BaseModel):
+    bucket: str
+    region: str = "us-east-1"
+    endpoint_url: str | None = None
+    access_key_id: str | None = None
+    secret_access_key: str | None = None
+
+
+class StorageSettings(BaseModel):
+    backend: StorageType = StorageType.FILESYSTEM
+    s3: S3StorageSettings | None = None
+
+
 class GooseBitSettings(BaseSettings):
-    model_config = SettingsConfigDict(env_prefix="GOOSEBIT_")
+    model_config = SettingsConfigDict(env_prefix="GOOSEBIT_", extra="ignore", env_nested_delimiter="__")
 
     port: int = 60053  # GOOSE
+    tenant: str = "DEFAULT"
+
+    poll_time: str = "00:01:00"
 
-    poll_time_default: str = "00:01:00"
-    poll_time_updating: str = "00:00:05"
-    poll_time_registration: str = "00:00:10"
+    max_concurrent_updates: int = 1000
+
+    device_auth: DeviceAuthSettings = DeviceAuthSettings()
 
     secret_key: Annotated[OctKey, BeforeValidator(OctKey.import_key)] = secrets.token_hex(16)
 
-    users: list[User] = []
+    plugins: list[str] = Field(default_factory=list)
 
     db_uri: str = f"sqlite:///{GOOSEBIT_ROOT_DIR.joinpath('db.sqlite3')}"
+    db_ssl_crt: Path | None = None
+
     artifacts_dir: Path = GOOSEBIT_ROOT_DIR.joinpath("artifacts")
 
+    storage: StorageSettings = StorageSettings()
+
     metrics: MetricsSettings = MetricsSettings()
 
     logging: dict = LOGGING_DEFAULT

goosebit/storage/__init__.py

@@ -0,0 +1,62 @@
+from pathlib import Path
+from typing import AsyncIterable
+
+from goosebit.settings import config
+from goosebit.settings.schema import GooseBitSettings, StorageType
+from goosebit.storage.base import StorageProtocol
+
+from .filesystem import FilesystemStorageBackend
+from .s3 import S3StorageBackend
+
+
+class GoosebitStorage:
+    def __init__(self, config: GooseBitSettings):
+        self.config = config
+        self._backend: StorageProtocol = self._create_backend()
+
+    def _create_backend(self) -> StorageProtocol:
+
+        if self.config.storage.backend == StorageType.FILESYSTEM:
+            return FilesystemStorageBackend(base_path=self.config.artifacts_dir)
+
+        elif self.config.storage.backend == StorageType.S3:
+            if self.config.storage.s3 is None:
+                return FilesystemStorageBackend(base_path=self.config.artifacts_dir)
+
+            s3_config = self.config.storage.s3
+            return S3StorageBackend(
+                bucket=s3_config.bucket,
+                region=s3_config.region,
+                endpoint_url=s3_config.endpoint_url,
+                access_key_id=s3_config.access_key_id,
+                secret_access_key=s3_config.secret_access_key,
+            )
+
+        else:
+            raise ValueError(f"Unknown storage backend type: {self.config.storage.backend}")
+
+    @property
+    def backend(self) -> StorageProtocol:
+        if self._backend is None:
+            raise RuntimeError("Storage backend not initialized. Use async context manager.")
+        return self._backend
+
+    async def store_file(self, source_path: Path, dest_path: Path) -> str:
+        return await self.backend.store_file(source_path, dest_path)
+
+    async def get_file_stream(self, uri: str) -> AsyncIterable[bytes]:
+        async for chunk in self.backend.get_file_stream(uri):  # type: ignore[attr-defined]
+            yield chunk
+
+    async def get_download_url(self, uri: str) -> str:
+        return await self.backend.get_download_url(uri)
+
+    def get_temp_dir(self) -> Path:
+        return self.backend.get_temp_dir()
+
+    async def delete_file(self, uri: str) -> bool:
+        return await self.backend.delete_file(uri)
+
+
+# Init the module-level storage instance
+storage = GoosebitStorage(config)

goosebit/storage/base.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+from typing import AsyncIterable, Protocol
+
+
+class StorageProtocol(Protocol):
+    async def store_file(self, source_path: Path, dest_path: Path) -> str: ...
+
+    async def get_file_stream(self, uri: str) -> AsyncIterable[bytes]: ...
+
+    async def get_download_url(self, uri: str) -> str: ...
+
+    async def delete_file(self, uri: str) -> bool: ...
+
+    def get_temp_dir(self) -> Path: ...