kumiho-construct 2026.4.28

//! Device management and pairing API handlers.

use super::{AppState, client_key_from_request};
use axum::{
    extract::{ConnectInfo, State},
    http::{HeaderMap, StatusCode, header},
    response::{IntoResponse, Json},
};
use chrono::{DateTime, Utc};
use parking_lot::Mutex;
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use tracing::{debug, error, info, warn};

/// Metadata about a paired device.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceInfo {
    pub id: String,
    pub name: Option<String>,
    pub device_type: Option<String>,
    pub paired_at: DateTime<Utc>,
    pub last_seen: DateTime<Utc>,
    pub ip_address: Option<String>,
}

/// Registry of paired devices backed by SQLite.
#[derive(Debug)]
pub struct DeviceRegistry {
    cache: Mutex<HashMap<String, DeviceInfo>>,
    db_path: PathBuf,
}

impl DeviceRegistry {
    /// Open (or create) the SQLite-backed device registry at
    /// `<workspace_dir>/devices.db`. Returns an error rather than panicking
    /// when the DB is locked or corrupt — a transient I/O failure at startup
    /// or on a request path must not take the gateway down.
    pub fn new(workspace_dir: &Path) -> anyhow::Result<Self> {
        use anyhow::Context;

        let db_path = workspace_dir.join("devices.db");
        let conn = Connection::open(&db_path)
            .with_context(|| format!("open device registry DB at {}", db_path.display()))?;
        conn.execute_batch(
            "CREATE TABLE IF NOT EXISTS devices (
                token_hash TEXT PRIMARY KEY,
                id TEXT NOT NULL,
                name TEXT,
                device_type TEXT,
                paired_at TEXT NOT NULL,
                last_seen TEXT NOT NULL,
                ip_address TEXT
            )",
        )
        .context("create devices table")?;

        let mut cache = HashMap::new();
        let mut stmt = conn
            .prepare("SELECT token_hash, id, name, device_type, paired_at, last_seen, ip_address FROM devices")
            .context("prepare device select")?;
        let rows = stmt
            .query_map([], |row| {
                let token_hash: String = row.get(0)?;
                let id: String = row.get(1)?;
                let name: Option<String> = row.get(2)?;
                let device_type: Option<String> = row.get(3)?;
                let paired_at_str: String = row.get(4)?;
                let last_seen_str: String = row.get(5)?;
                let ip_address: Option<String> = row.get(6)?;
                let paired_at = DateTime::parse_from_rfc3339(&paired_at_str)
                    .map(|dt| dt.with_timezone(&Utc))
                    .unwrap_or_else(|_| Utc::now());
                let last_seen = DateTime::parse_from_rfc3339(&last_seen_str)
                    .map(|dt| dt.with_timezone(&Utc))
                    .unwrap_or_else(|_| Utc::now());
                Ok((
                    token_hash,
                    DeviceInfo {
                        id,
                        name,
                        device_type,
                        paired_at,
                        last_seen,
                        ip_address,
                    },
                ))
            })
            .context("query devices")?;
        for (hash, info) in rows.flatten() {
            cache.insert(hash, info);
        }

        Ok(Self {
            cache: Mutex::new(cache),
            db_path,
        })
    }

    fn open_db(&self) -> anyhow::Result<Connection> {
        use anyhow::Context;
        Connection::open(&self.db_path)
            .with_context(|| format!("open device registry DB at {}", self.db_path.display()))
    }

    pub fn register(&self, token_hash: String, info: DeviceInfo) -> anyhow::Result<()> {
        use anyhow::Context;
        let conn = self.open_db()?;
        let device_id = info.id.clone();
        conn.execute(
            "INSERT OR REPLACE INTO devices (token_hash, id, name, device_type, paired_at, last_seen, ip_address) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            rusqlite::params![
                token_hash,
                info.id,
                info.name,
                info.device_type,
                info.paired_at.to_rfc3339(),
                info.last_seen.to_rfc3339(),
                info.ip_address,
            ],
        )
        .context("insert device row")?;
        let hash_prefix: String = token_hash.chars().take(8).collect();
        self.cache.lock().insert(token_hash, info);
        info!(device_id = %device_id, token_hash_prefix = %hash_prefix, "device registered in SQLite");
        Ok(())
    }

    pub fn list(&self) -> Vec<DeviceInfo> {
        // Persistence failures here degrade to an empty list rather than
        // panicking: the caller is a request handler, not startup code.
        let conn = match self.open_db() {
            Ok(c) => c,
            Err(e) => {
                warn!(error = %e, "device registry list: open_db failed — returning empty list");
                return Vec::new();
            }
        };
        let mut stmt = match conn.prepare(
            "SELECT token_hash, id, name, device_type, paired_at, last_seen, ip_address FROM devices",
        ) {
            Ok(s) => s,
            Err(e) => {
                warn!(error = %e, "device registry list: prepare failed — returning empty list");
                return Vec::new();
            }
        };
        let rows = match stmt.query_map([], |row| {
            let id: String = row.get(1)?;
            let name: Option<String> = row.get(2)?;
            let device_type: Option<String> = row.get(3)?;
            let paired_at_str: String = row.get(4)?;
            let last_seen_str: String = row.get(5)?;
            let ip_address: Option<String> = row.get(6)?;
            let paired_at = DateTime::parse_from_rfc3339(&paired_at_str)
                .map(|dt| dt.with_timezone(&Utc))
                .unwrap_or_else(|_| Utc::now());
            let last_seen = DateTime::parse_from_rfc3339(&last_seen_str)
                .map(|dt| dt.with_timezone(&Utc))
                .unwrap_or_else(|_| Utc::now());
            Ok(DeviceInfo {
                id,
                name,
                device_type,
                paired_at,
                last_seen,
                ip_address,
            })
        }) {
            Ok(r) => r,
            Err(e) => {
                warn!(error = %e, "device registry list: query_map failed — returning empty list");
                return Vec::new();
            }
        };
        rows.filter_map(|r| r.ok()).collect()
    }

    pub fn revoke(&self, device_id: &str) -> bool {
        let conn = match self.open_db() {
            Ok(c) => c,
            Err(e) => {
                warn!(error = %e, "device registry revoke: open_db failed");
                return false;
            }
        };
        let deleted = conn
            .execute(
                "DELETE FROM devices WHERE id = ?1",
                rusqlite::params![device_id],
            )
            .unwrap_or(0);
        if deleted > 0 {
            let mut cache = self.cache.lock();
            let key = cache
                .iter()
                .find(|(_, v)| v.id == device_id)
                .map(|(k, _)| k.clone());
            if let Some(key) = key {
                cache.remove(&key);
            }
            true
        } else {
            false
        }
    }

    pub fn update_last_seen(&self, token_hash: &str) {
        let now = Utc::now();
        if let Ok(conn) = self.open_db() {
            let _ = conn.execute(
                "UPDATE devices SET last_seen = ?1 WHERE token_hash = ?2",
                rusqlite::params![now.to_rfc3339(), token_hash],
            );
        }
        if let Some(device) = self.cache.lock().get_mut(token_hash) {
            device.last_seen = now;
        }
    }

    pub fn device_count(&self) -> usize {
        self.cache.lock().len()
    }
}

/// Store for pending pairing requests.
#[derive(Debug)]
pub struct PairingStore {
    pending: Mutex<Vec<PendingPairing>>,
    max_pending: usize,
}

#[derive(Debug, Clone, Serialize)]
struct PendingPairing {
    code: String,
    created_at: DateTime<Utc>,
    expires_at: DateTime<Utc>,
    client_ip: Option<String>,
    attempts: u32,
}

impl PairingStore {
    pub fn new(max_pending: usize) -> Self {
        Self {
            pending: Mutex::new(Vec::new()),
            max_pending,
        }
    }

    pub fn pending_count(&self) -> usize {
        let mut pending = self.pending.lock();
        pending.retain(|p| p.expires_at > Utc::now());
        pending.len()
    }
}

fn extract_bearer(headers: &HeaderMap) -> Option<&str> {
    headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|auth| auth.strip_prefix("Bearer "))
}

fn require_auth(state: &AppState, headers: &HeaderMap) -> Result<(), (StatusCode, &'static str)> {
    if state.pairing.require_pairing() {
        let token = extract_bearer(headers).unwrap_or("");
        if !state.pairing.is_authenticated(token) {
            return Err((StatusCode::UNAUTHORIZED, "Unauthorized"));
        }
    }
    Ok(())
}

/// POST /api/pairing/initiate — initiate a new pairing session
pub async fn initiate_pairing(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> impl IntoResponse {
    if let Err(e) = require_auth(&state, &headers) {
        warn!("initiate_pairing: unauthorized request");
        return e.into_response();
    }

    info!("initiate_pairing: generating new pairing code");
    match state.pairing.generate_new_pairing_code() {
        Some(code) => {
            let code_prefix: String = code.chars().take(2).collect();
            info!(code_prefix = %code_prefix, "initiate_pairing: code generated");
            Json(serde_json::json!({
                "pairing_code": code,
                "message": "New pairing code generated"
            }))
            .into_response()
        }
        None => {
            warn!("initiate_pairing: pairing disabled or unavailable");
            (
                StatusCode::SERVICE_UNAVAILABLE,
                "Pairing is disabled or not available",
            )
                .into_response()
        }
    }
}

/// POST /api/pair — submit pairing code (for new device pairing)
pub async fn submit_pairing_enhanced(
    State(state): State<AppState>,
    ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
    headers: HeaderMap,
    Json(body): Json<serde_json::Value>,
) -> impl IntoResponse {
    let code = body["code"].as_str().unwrap_or("");
    let device_name = body["device_name"].as_str().map(String::from);
    let device_type = body["device_type"].as_str().map(String::from);

    // Use the shared client-key helper. This ignores unsolicited
    // `X-Forwarded-For` values unless `trust_forwarded_headers` is set, so
    // attackers cannot bypass the per-client lockout by rotating the header.
    let client_id =
        client_key_from_request(Some(peer_addr), &headers, state.trust_forwarded_headers);

    info!(
        client_id = %client_id,
        code_len = code.len(),
        device_name = ?device_name,
        device_type = ?device_type,
        "submit_pairing_enhanced: received pair request"
    );

    match state.pairing.try_pair(code, &client_id).await {
        Ok(Some(token)) => {
            // Register the new device
            let token_hash = {
                use sha2::{Digest, Sha256};
                let hash = Sha256::digest(token.as_bytes());
                hex::encode(hash)
            };
            let hash_prefix: String = token_hash.chars().take(8).collect();
            info!(
                client_id = %client_id,
                token_hash_prefix = %hash_prefix,
                "submit_pairing_enhanced: pairing succeeded, registering device"
            );
            if let Some(ref registry) = state.device_registry {
                if let Err(e) = registry.register(
                    token_hash,
                    DeviceInfo {
                        id: uuid::Uuid::new_v4().to_string(),
                        name: device_name,
                        device_type,
                        paired_at: Utc::now(),
                        last_seen: Utc::now(),
                        ip_address: Some(client_id.clone()),
                    },
                ) {
                    error!(
                        client_id = %client_id,
                        error = %e,
                        "submit_pairing_enhanced: device registry insert failed"
                    );
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        "Pairing succeeded but device registration failed",
                    )
                        .into_response();
                }
            } else {
                debug!("submit_pairing_enhanced: no device_registry configured; skipping persist");
            }
            Json(serde_json::json!({
                "token": token,
                "message": "Pairing successful"
            }))
            .into_response()
        }
        Ok(None) => {
            warn!(client_id = %client_id, "submit_pairing_enhanced: invalid or expired code");
            (StatusCode::BAD_REQUEST, "Invalid or expired pairing code").into_response()
        }
        Err(lockout_secs) => {
            warn!(
                client_id = %client_id,
                lockout_secs,
                "submit_pairing_enhanced: client locked out"
            );
            (
                StatusCode::TOO_MANY_REQUESTS,
                format!("Too many attempts. Locked out for {lockout_secs}s"),
            )
                .into_response()
        }
    }
}

/// GET /api/devices — list paired devices
pub async fn list_devices(State(state): State<AppState>, headers: HeaderMap) -> impl IntoResponse {
    if let Err(e) = require_auth(&state, &headers) {
        return e.into_response();
    }

    let devices = state
        .device_registry
        .as_ref()
        .map(|r| r.list())
        .unwrap_or_default();

    let count = devices.len();
    Json(serde_json::json!({
        "devices": devices,
        "count": count
    }))
    .into_response()
}

/// DELETE /api/devices/{id} — revoke a paired device
pub async fn revoke_device(
    State(state): State<AppState>,
    headers: HeaderMap,
    axum::extract::Path(device_id): axum::extract::Path<String>,
) -> impl IntoResponse {
    if let Err(e) = require_auth(&state, &headers) {
        return e.into_response();
    }

    let revoked = state
        .device_registry
        .as_ref()
        .map(|r| r.revoke(&device_id))
        .unwrap_or(false);

    if revoked {
        Json(serde_json::json!({
            "message": "Device revoked",
            "device_id": device_id
        }))
        .into_response()
    } else {
        (StatusCode::NOT_FOUND, "Device not found").into_response()
    }
}

/// POST /api/devices/{id}/token/rotate — rotate a device's token
pub async fn rotate_token(
    State(state): State<AppState>,
    headers: HeaderMap,
    axum::extract::Path(device_id): axum::extract::Path<String>,
) -> impl IntoResponse {
    if let Err(e) = require_auth(&state, &headers) {
        return e.into_response();
    }

    // Generate a new pairing code for re-pairing
    match state.pairing.generate_new_pairing_code() {
        Some(code) => Json(serde_json::json!({
            "device_id": device_id,
            "pairing_code": code,
            "message": "Use this code to re-pair the device"
        }))
        .into_response(),
        None => (
            StatusCode::SERVICE_UNAVAILABLE,
            "Cannot generate new pairing code",
        )
            .into_response(),
    }
}