""" Device registration endpoint. """ import asyncio import copy import json import secrets from base64 import b64encode from datetime import datetime, timezone from pathlib import Path from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, status from pydantic import BaseModel from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession from app.core.database import get_db from app.models.device import Device from app.models.settings import Settings from app.utils.ssh_keys import sync_authorized_keys router = APIRouter() # Load default config from JSON file DEFAULT_CONFIG_PATH = Path(__file__).parent.parent.parent / "default_config.json" with open(DEFAULT_CONFIG_PATH, "r") as f: DEFAULT_CONFIG = json.load(f) class RegistrationRequest(BaseModel): """Device registration request.""" device_id: str # MAC address ssh_public_key: str | None = None class RegistrationResponse(BaseModel): """Device registration response.""" device_token: str device_password: str def _generate_token() -> str: """Generate 32-byte base64 token.""" return b64encode(secrets.token_bytes(32)).decode("ascii") def _generate_password() -> str: """Generate 8-digit password.""" n = secrets.randbelow(10**8) return f"{n:08d}" @router.post("/registration", response_model=RegistrationResponse, status_code=201) async def register_device( data: RegistrationRequest, db: Annotated[AsyncSession, Depends(get_db)], ): """ Register new device or return existing credentials. Requires auto_registration to be enabled in settings. """ mac_address = data.device_id.lower().strip() if not mac_address: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Missing device_id", ) # Check if device already exists result = await db.execute(select(Device).where(Device.mac_address == mac_address)) device = result.scalar_one_or_none() if device: # Return existing credentials if not device.device_token or not device.device_password: # Re-generate if missing device.device_token = _generate_token() device.device_password = _generate_password() await db.commit() return RegistrationResponse( device_token=device.device_token, device_password=device.device_password, ) # Check auto-registration setting settings_result = await db.execute( select(Settings).where(Settings.key == "auto_registration") ) auto_reg_setting = settings_result.scalar_one_or_none() if not auto_reg_setting or not auto_reg_setting.value.get("enabled", False): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Registration disabled. Contact administrator.", ) # Create new device with default config # Deep copy to avoid modifying the global default device_config = copy.deepcopy(DEFAULT_CONFIG) # Add SSH public key if provided if data.ssh_public_key: device_config["ssh_public_key"] = data.ssh_public_key device = Device( mac_address=mac_address, organization_id=None, # Unassigned status="online", config=device_config, device_token=_generate_token(), device_password=_generate_password(), ) db.add(device) await db.flush() # Update last_device_at auto_reg_setting.value["last_device_at"] = datetime.now(timezone.utc).isoformat() await db.commit() await db.refresh(device) print(f"[REGISTRATION] device={mac_address} simple_id={device.simple_id}") # Sync SSH keys to authorized_keys (background task) if data.ssh_public_key: asyncio.create_task(sync_authorized_keys()) return RegistrationResponse( device_token=device.device_token, device_password=device.device_password, )