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

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