""" Device registration endpoint. """ import asyncio import copy import json import secrets import subprocess 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 async_session_maker, get_db from app.models.device import Device from app.models.settings import Settings 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}" async def _sync_authorized_keys(): """Sync all device SSH keys to /home/tunnel/.ssh/authorized_keys""" try: async with async_session_maker() as session: result = await session.execute(select(Device)) devices = result.scalars().all() keys = [] for device in devices: if device.config and 'ssh_public_key' in device.config: ssh_key = device.config['ssh_public_key'].strip() if ssh_key: keys.append(f"{ssh_key} # {device.mac_address}") authorized_keys_content = "\n".join(keys) + "\n" if keys else "" # Write using sudo subprocess.run( ["sudo", "tee", "/home/tunnel/.ssh/authorized_keys"], input=authorized_keys_content.encode(), stdout=subprocess.DEVNULL, check=True ) subprocess.run( ["sudo", "chmod", "600", "/home/tunnel/.ssh/authorized_keys"], check=True ) subprocess.run( ["sudo", "chown", "tunnel:tunnel", "/home/tunnel/.ssh/authorized_keys"], check=True ) print(f"[SSH] Synced {len(keys)} keys to authorized_keys") except Exception as e: print(f"[SSH] Failed to sync authorized_keys: {e}") @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, )