mybeacon-backend/backend/app/api/v1/registration.py

"""
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:
        # Update SSH key if provided (device may have regenerated keys)
        ssh_key_updated = False
        if data.ssh_public_key:
            new_key = data.ssh_public_key.strip()
            old_key = (device.config or {}).get("ssh_public_key", "").strip()
            if new_key != old_key:
                # Update config with new key (preserve other settings)
                new_config = {**(device.config or {}), "ssh_public_key": new_key}
                device.config = new_config
                ssh_key_updated = True
                print(f"[REGISTRATION] Updated SSH key for device={mac_address}")

        # Re-generate credentials if missing
        if not device.device_token or not device.device_password:
            device.device_token = _generate_token()
            device.device_password = _generate_password()

        if ssh_key_updated or not device.device_token:
            await db.commit()

        # Sync SSH keys if updated
        if ssh_key_updated:
            asyncio.create_task(sync_authorized_keys())

        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,
    )