epaper-display/LORAWEB_NEU/backend/api/src/users.rs

use axum::{
    extract::{Path, State},
    http::StatusCode,
    routing::{delete, get, post, put},
    Json, Router,
};
use argon2::{Argon2, PasswordHasher, password_hash::SaltString};
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use chrono::Utc;
use crate::AppState;

#[derive(Debug, Serialize, sqlx::FromRow)]
struct User {
    id: i64,
    username: String,
    role: String,
    totp_enabled: i32,
    created_at: String,
}

#[derive(Deserialize)]
struct CreateUserRequest {
    username: String,
    password: String,
    role: Option<String>,
}

#[derive(Deserialize)]
struct UpdateUserRequest {
    username: Option<String>,
    password: Option<String>,
    role: Option<String>,
}

#[derive(Serialize)]
struct TotpSetupResponse {
    secret: String,
    uri: String,
}

type ApiError = (StatusCode, Json<serde_json::Value>);

fn internal_err(msg: &str) -> ApiError {
    (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": msg })))
}

fn not_found(msg: &str) -> ApiError {
    (StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": msg })))
}

fn bad_request(msg: &str) -> ApiError {
    (StatusCode::BAD_REQUEST, Json(serde_json::json!({ "error": msg })))
}

fn conflict(msg: &str) -> ApiError {
    (StatusCode::CONFLICT, Json(serde_json::json!({ "error": msg })))
}

pub fn routes() -> Router<AppState> {
    Router::new()
        .route("/users", get(list_users).post(create_user))
        .route("/users/{id}", put(update_user).delete(delete_user))
        .route("/users/{id}/totp", post(enable_totp).delete(disable_totp))
}

async fn list_users(
    State(state): State<AppState>,
) -> Result<Json<Vec<User>>, ApiError> {
    let users = sqlx::query_as::<_, User>(
        "SELECT id, username, role, totp_enabled, created_at FROM users ORDER BY created_at ASC"
    )
    .fetch_all(&state.pool)
    .await
    .map_err(|_| internal_err("Database error"))?;
    Ok(Json(users))
}

async fn create_user(
    State(state): State<AppState>,
    Json(req): Json<CreateUserRequest>,
) -> Result<(StatusCode, Json<User>), ApiError> {
    if req.username.trim().is_empty() {
        return Err(bad_request("Username is required"));
    }
    if req.password.is_empty() {
        return Err(bad_request("Password is required"));
    }

    let role = req.role.unwrap_or_else(|| "user".into());
    let salt = SaltString::generate(&mut OsRng);
    let hash = Argon2::default()
        .hash_password(req.password.as_bytes(), &salt)
        .map_err(|_| internal_err("Failed to hash password"))?
        .to_string();
    let now = Utc::now().to_rfc3339();

    let result = sqlx::query(
        "INSERT INTO users (username, password_hash, role, totp_enabled, created_at) VALUES (?, ?, ?, 0, ?)"
    )
    .bind(&req.username)
    .bind(&hash)
    .bind(&role)
    .bind(&now)
    .execute(&state.pool)
    .await;

    let insert_result = match result {
        Ok(r) => r,
        Err(e) => {
            let msg = e.to_string();
            if msg.contains("UNIQUE") || msg.contains("unique") {
                return Err(conflict("Username already exists"));
            }
            return Err(internal_err("Failed to create user"));
        }
    };

    let user_id = insert_result.last_insert_rowid();
    let user = sqlx::query_as::<_, User>(
        "SELECT id, username, role, totp_enabled, created_at FROM users WHERE id = ?"
    )
    .bind(user_id)
    .fetch_one(&state.pool)
    .await
    .map_err(|_| internal_err("Failed to fetch created user"))?;

    Ok((StatusCode::CREATED, Json(user)))
}

async fn update_user(
    State(state): State<AppState>,
    Path(id): Path<i64>,
    Json(req): Json<UpdateUserRequest>,
) -> Result<Json<User>, ApiError> {
    let existing = sqlx::query_as::<_, User>(
        "SELECT id, username, role, totp_enabled, created_at FROM users WHERE id = ?"
    )
    .bind(id)
    .fetch_optional(&state.pool)
    .await
    .map_err(|_| internal_err("Database error"))?
    .ok_or_else(|| not_found("User not found"))?;

    let username = req.username.unwrap_or(existing.username);
    let role = req.role.unwrap_or(existing.role);

    if let Some(password) = req.password {
        if password.is_empty() {
            return Err(bad_request("Password cannot be empty"));
        }
        let salt = SaltString::generate(&mut OsRng);
        let hash = Argon2::default()
            .hash_password(password.as_bytes(), &salt)
            .map_err(|_| internal_err("Failed to hash password"))?
            .to_string();

        sqlx::query("UPDATE users SET username=?, role=?, password_hash=? WHERE id=?")
            .bind(&username)
            .bind(&role)
            .bind(&hash)
            .bind(id)
            .execute(&state.pool)
            .await
            .map_err(|_| internal_err("Failed to update user"))?;
    } else {
        sqlx::query("UPDATE users SET username=?, role=? WHERE id=?")
            .bind(&username)
            .bind(&role)
            .bind(id)
            .execute(&state.pool)
            .await
            .map_err(|_| internal_err("Failed to update user"))?;
    }

    let user = sqlx::query_as::<_, User>(
        "SELECT id, username, role, totp_enabled, created_at FROM users WHERE id = ?"
    )
    .bind(id)
    .fetch_one(&state.pool)
    .await
    .map_err(|_| internal_err("Failed to fetch updated user"))?;

    Ok(Json(user))
}

async fn delete_user(
    State(state): State<AppState>,
    Path(id): Path<i64>,
) -> Result<StatusCode, ApiError> {
    let rows_affected = sqlx::query("DELETE FROM users WHERE id=?")
        .bind(id)
        .execute(&state.pool)
        .await
        .map_err(|_| internal_err("Database error"))?
        .rows_affected();

    if rows_affected == 0 {
        return Err(not_found("User not found"));
    }

    Ok(StatusCode::NO_CONTENT)
}

async fn enable_totp(
    State(state): State<AppState>,
    Path(id): Path<i64>,
) -> Result<Json<TotpSetupResponse>, ApiError> {
    // Verify user exists and get username for the TOTP URI
    let row: Option<(String,)> = sqlx::query_as("SELECT username FROM users WHERE id = ?")
        .bind(id)
        .fetch_optional(&state.pool)
        .await
        .map_err(|_| internal_err("Database error"))?;

    let (username,) = row.ok_or_else(|| not_found("User not found"))?;

    let secret = totp_rs::Secret::generate_secret();
    let secret_encoded = secret.to_encoded().to_string();

    let secret_bytes = totp_rs::Secret::Encoded(secret_encoded.clone())
        .to_bytes()
        .map_err(|_| internal_err("Failed to decode TOTP secret"))?;

    let totp = totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        6,
        1,
        30,
        secret_bytes,
    )
    .map_err(|_| internal_err("Failed to create TOTP"))?;

    let uri = totp.get_url(&username, "LoRaWEB");

    sqlx::query("UPDATE users SET totp_secret=?, totp_enabled=1 WHERE id=?")
        .bind(&secret_encoded)
        .bind(id)
        .execute(&state.pool)
        .await
        .map_err(|_| internal_err("Failed to save TOTP secret"))?;

    Ok(Json(TotpSetupResponse {
        secret: secret_encoded,
        uri,
    }))
}

async fn disable_totp(
    State(state): State<AppState>,
    Path(id): Path<i64>,
) -> Result<StatusCode, ApiError> {
    let rows_affected = sqlx::query("UPDATE users SET totp_secret=NULL, totp_enabled=0 WHERE id=?")
        .bind(id)
        .execute(&state.pool)
        .await
        .map_err(|_| internal_err("Database error"))?
        .rows_affected();

    if rows_affected == 0 {
        return Err(not_found("User not found"));
    }

    Ok(StatusCode::NO_CONTENT)
}