""" Device registration endpoint. """ import secrets from base64 import b64encode from datetime import datetime, timezone 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 router = APIRouter() 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 device = Device( mac_address=mac_address, organization_id=None, # Unassigned status="online", config={"ssh_public_key": data.ssh_public_key} if data.ssh_public_key else {}, 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}") return RegistrationResponse( device_token=device.device_token, device_password=device.device_password, )