athena_rs 3.3.0

use chrono::{DateTime, Utc};
use serde::Serialize;
use serde_json::{Value, json};
use sqlx::postgres::{PgPool, PgRow};
use sqlx::{Postgres, Row, Transaction};
use thiserror::Error;
use uuid::Uuid;

#[derive(Debug, Clone, Serialize)]
pub struct ApiKeyRecord {
    pub id: String,
    pub public_id: String,
    pub name: String,
    pub description: Option<String>,
    pub client_name: Option<String>,
    pub expires_at: Option<DateTime<Utc>>,
    pub is_active: bool,
    pub last_used_at: Option<DateTime<Utc>>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub rights: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct ApiKeySecretRecord {
    pub record: ApiKeyRecord,
    pub internal_id: i64,
    pub key_salt: String,
    pub key_hash: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct ApiKeyRightRecord {
    pub id: String,
    pub name: String,
    pub description: Option<String>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ApiKeyGlobalConfig {
    pub enforce_api_keys: bool,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ApiKeyClientConfig {
    pub client_name: String,
    pub enforce_api_keys: bool,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone)]
pub struct CreateApiKeyParams {
    pub public_id: String,
    pub name: String,
    pub description: Option<String>,
    pub client_name: Option<String>,
    pub key_salt: String,
    pub key_hash: String,
    pub expires_at: Option<DateTime<Utc>>,
    pub rights: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct SaveApiKeyParams {
    pub id: Uuid,
    pub name: String,
    pub description: Option<String>,
    pub client_name: Option<String>,
    pub expires_at: Option<DateTime<Utc>>,
    pub is_active: bool,
    pub rights: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct AuthAttemptLogEntry {
    pub request_id: String,
    pub api_key_id: Option<String>,
    pub internal_api_key_id: Option<i64>,
    pub api_key_public_id: Option<String>,
    pub client_name: Option<String>,
    pub method: String,
    pub path: String,
    pub presented_key_hash: Option<String>,
    pub presented_key_salt: Option<String>,
    pub required_rights: Vec<String>,
    pub granted_rights: Vec<String>,
    pub authenticated: bool,
    pub authorized: bool,
    pub enforced: bool,
    pub failure_reason: Option<String>,
    pub remote_addr: Option<String>,
    pub user_agent: Option<String>,
    pub time: i64,
}

#[derive(Debug, Error)]
pub enum ApiKeyStoreError {
    #[error("database error: {0}")]
    Database(#[from] sqlx::Error),
    #[error("missing rights: {0:?}")]
    MissingRights(Vec<String>),
}

const API_KEY_SELECT_FIELDS: &str = r#"
    k.api_keys_id AS id,
    k.public_id,
    k.name,
    k.description,
    k.client_name,
    k.expires_at,
    k.is_active,
    k.last_used_at,
    k.created_at,
    k.updated_at,
    COALESCE(
        array_agg(r.name ORDER BY r.name) FILTER (WHERE r.name IS NOT NULL),
        ARRAY[]::text[]
    ) AS rights
"#;

fn map_api_key_record(row: &PgRow) -> Result<ApiKeyRecord, sqlx::Error> {
    Ok(ApiKeyRecord {
        id: row.try_get::<Uuid, _>("id")?.to_string(),
        public_id: row.try_get("public_id")?,
        name: row.try_get("name")?,
        description: row.try_get("description")?,
        client_name: row.try_get("client_name")?,
        expires_at: row.try_get("expires_at")?,
        is_active: row.try_get("is_active")?,
        last_used_at: row.try_get("last_used_at")?,
        created_at: row.try_get("created_at")?,
        updated_at: row.try_get("updated_at")?,
        rights: row.try_get("rights")?,
    })
}

async fn ensure_rights_exist(
    tx: &mut Transaction<'_, Postgres>,
    rights: &[String],
) -> Result<(), ApiKeyStoreError> {
    if rights.is_empty() {
        return Ok(());
    }

    let existing: Vec<String> = sqlx::query_scalar(
        r#"
        SELECT name
        FROM api_key_rights
        WHERE name = ANY($1)
        "#,
    )
    .bind(rights)
    .fetch_all(&mut **tx)
    .await?;

    let missing: Vec<String> = rights
        .iter()
        .filter(|right| !existing.iter().any(|candidate| candidate == *right))
        .cloned()
        .collect();
    if !missing.is_empty() {
        return Err(ApiKeyStoreError::MissingRights(missing));
    }

    Ok(())
}

pub async fn list_api_keys(pool: &PgPool) -> Result<Vec<ApiKeyRecord>, sqlx::Error> {
    let rows: Vec<PgRow> = sqlx::query(&format!(
        r#"
        SELECT {fields}
        FROM api_keys k
        LEFT JOIN api_key_right_grants g ON g.api_key_id = k.id
        LEFT JOIN api_key_rights r ON r.id = g.right_id
        GROUP BY k.id, k.api_keys_id
        ORDER BY k.created_at DESC
        "#,
        fields = API_KEY_SELECT_FIELDS
    ))
    .fetch_all(pool)
    .await?;

    rows.iter().map(map_api_key_record).collect()
}

pub async fn get_api_key(pool: &PgPool, id: Uuid) -> Result<Option<ApiKeyRecord>, sqlx::Error> {
    let row: Option<PgRow> = sqlx::query(&format!(
        r#"
        SELECT {fields}
        FROM api_keys k
        LEFT JOIN api_key_right_grants g ON g.api_key_id = k.id
        LEFT JOIN api_key_rights r ON r.id = g.right_id
        WHERE k.api_keys_id = $1
        GROUP BY k.id, k.api_keys_id
        "#,
        fields = API_KEY_SELECT_FIELDS
    ))
    .bind(id)
    .fetch_optional(pool)
    .await?;

    row.as_ref().map(map_api_key_record).transpose()
}

pub async fn get_api_key_by_public_id(
    pool: &PgPool,
    public_id: &str,
) -> Result<Option<ApiKeySecretRecord>, sqlx::Error> {
    let row: Option<PgRow> = sqlx::query(&format!(
        r#"
        SELECT
            k.id AS internal_id,
            {fields},
            k.key_salt,
            k.key_hash
        FROM api_keys k
        LEFT JOIN api_key_right_grants g ON g.api_key_id = k.id
        LEFT JOIN api_key_rights r ON r.id = g.right_id
        WHERE k.public_id = $1
        GROUP BY k.id, k.api_keys_id
        "#,
        fields = API_KEY_SELECT_FIELDS
    ))
    .bind(public_id)
    .fetch_optional(pool)
    .await?;

    row.as_ref()
        .map(|value| {
            Ok(ApiKeySecretRecord {
                record: map_api_key_record(value)?,
                internal_id: value.try_get("internal_id")?,
                key_salt: value.try_get("key_salt")?,
                key_hash: value.try_get("key_hash")?,
            })
        })
        .transpose()
}

pub async fn create_api_key(
    pool: &PgPool,
    params: CreateApiKeyParams,
) -> Result<ApiKeyRecord, ApiKeyStoreError> {
    let api_keys_id: Uuid = Uuid::new_v4();
    let mut tx = pool.begin().await?;

    ensure_rights_exist(&mut tx, &params.rights).await?;

    let row: PgRow = sqlx::query(
        r#"
        INSERT INTO api_keys
            (api_keys_id, public_id, name, description, client_name, key_salt, key_hash, expires_at, is_active)
        VALUES
            ($1, $2, $3, $4, $5, $6, $7, $8, true)
        RETURNING id
        "#,
    )
    .bind(api_keys_id)
    .bind(&params.public_id)
    .bind(&params.name)
    .bind(&params.description)
    .bind(&params.client_name)
    .bind(&params.key_salt)
    .bind(&params.key_hash)
    .bind(params.expires_at)
    .fetch_one(&mut *tx)
    .await?;

    let internal_id: i64 = row.try_get("id")?;

    if !params.rights.is_empty() {
        sqlx::query(
            r#"
            INSERT INTO api_key_right_grants (api_key_id, right_id)
            SELECT $1, r.id
            FROM api_key_rights r
            WHERE r.name = ANY($2)
            "#,
        )
        .bind(internal_id)
        .bind(&params.rights)
        .execute(&mut *tx)
        .await?;
    }

    tx.commit().await?;
    get_api_key(pool, api_keys_id)
        .await?
        .ok_or_else(|| sqlx::Error::RowNotFound.into())
}

pub async fn save_api_key(
    pool: &PgPool,
    params: SaveApiKeyParams,
) -> Result<Option<ApiKeyRecord>, ApiKeyStoreError> {
    let mut tx = pool.begin().await?;

    ensure_rights_exist(&mut tx, &params.rights).await?;

    let result = sqlx::query(
        r#"
        UPDATE api_keys
        SET
            name = $2,
            description = $3,
            client_name = $4,
            expires_at = $5,
            is_active = $6,
            updated_at = now()
        WHERE api_keys_id = $1
        "#,
    )
    .bind(params.id)
    .bind(&params.name)
    .bind(&params.description)
    .bind(&params.client_name)
    .bind(params.expires_at)
    .bind(params.is_active)
    .execute(&mut *tx)
    .await?;

    if result.rows_affected() == 0 {
        tx.rollback().await?;
        return Ok(None);
    }

    let internal_id: i64 = sqlx::query_scalar(
        r#"
        SELECT id
        FROM api_keys
        WHERE api_keys_id = $1
        "#,
    )
    .bind(params.id)
    .fetch_one(&mut *tx)
    .await?;

    sqlx::query(
        r#"
        DELETE FROM api_key_right_grants
        WHERE api_key_id = $1
        "#,
    )
    .bind(internal_id)
    .execute(&mut *tx)
    .await?;

    if !params.rights.is_empty() {
        sqlx::query(
            r#"
            INSERT INTO api_key_right_grants (api_key_id, right_id)
            SELECT $1, r.id
            FROM api_key_rights r
            WHERE r.name = ANY($2)
            "#,
        )
        .bind(internal_id)
        .bind(&params.rights)
        .execute(&mut *tx)
        .await?;
    }

    tx.commit().await?;
    get_api_key(pool, params.id)
        .await
        .map_err(ApiKeyStoreError::from)
}

pub async fn delete_api_key(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
    let result = sqlx::query(
        r#"
        DELETE FROM api_keys
        WHERE api_keys_id = $1
        "#,
    )
    .bind(id)
    .execute(pool)
    .await?;
    Ok(result.rows_affected() > 0)
}

pub async fn touch_api_key_last_used(pool: &PgPool, id: Uuid) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
        UPDATE api_keys
        SET
            last_used_at = now(),
            updated_at = now()
        WHERE api_keys_id = $1
        "#,
    )
    .bind(id)
    .execute(pool)
    .await?;
    Ok(())
}

pub async fn list_api_key_rights(pool: &PgPool) -> Result<Vec<ApiKeyRightRecord>, sqlx::Error> {
    let rows: Vec<PgRow> = sqlx::query(
        r#"
        SELECT api_key_rights_id AS id, name, description, created_at, updated_at
        FROM api_key_rights
        ORDER BY name
        "#,
    )
    .fetch_all(pool)
    .await?;

    rows.into_iter()
        .map(|row| {
            Ok(ApiKeyRightRecord {
                id: row.try_get::<Uuid, _>("id")?.to_string(),
                name: row.try_get("name")?,
                description: row.try_get("description")?,
                created_at: row.try_get("created_at")?,
                updated_at: row.try_get("updated_at")?,
            })
        })
        .collect()
}

pub async fn get_api_key_right(
    pool: &PgPool,
    id: Uuid,
) -> Result<Option<ApiKeyRightRecord>, sqlx::Error> {
    let row: Option<PgRow> = sqlx::query(
        r#"
        SELECT api_key_rights_id AS id, name, description, created_at, updated_at
        FROM api_key_rights
        WHERE api_key_rights_id = $1
        "#,
    )
    .bind(id)
    .fetch_optional(pool)
    .await?;

    row.map(|row| {
        Ok(ApiKeyRightRecord {
            id: row.try_get::<Uuid, _>("id")?.to_string(),
            name: row.try_get("name")?,
            description: row.try_get("description")?,
            created_at: row.try_get("created_at")?,
            updated_at: row.try_get("updated_at")?,
        })
    })
    .transpose()
}

pub async fn create_api_key_right(
    pool: &PgPool,
    name: &str,
    description: Option<&str>,
) -> Result<ApiKeyRightRecord, sqlx::Error> {
    let api_key_rights_id: Uuid = Uuid::new_v4();
    sqlx::query(
        r#"
        INSERT INTO api_key_rights (api_key_rights_id, name, description)
        VALUES ($1, $2, $3)
        "#,
    )
    .bind(api_key_rights_id)
    .bind(name)
    .bind(description)
    .execute(pool)
    .await?;

    get_api_key_right(pool, api_key_rights_id)
        .await?
        .ok_or(sqlx::Error::RowNotFound)
}

pub async fn update_api_key_right(
    pool: &PgPool,
    id: Uuid,
    name: &str,
    description: Option<&str>,
) -> Result<Option<ApiKeyRightRecord>, sqlx::Error> {
    let result = sqlx::query(
        r#"
        UPDATE api_key_rights
        SET
            name = $2,
            description = $3,
            updated_at = now()
        WHERE api_key_rights_id = $1
        "#,
    )
    .bind(id)
    .bind(name)
    .bind(description)
    .execute(pool)
    .await?;

    if result.rows_affected() == 0 {
        return Ok(None);
    }

    get_api_key_right(pool, id).await
}

pub async fn delete_api_key_right(pool: &PgPool, id: Uuid) -> Result<bool, sqlx::Error> {
    let result = sqlx::query(
        r#"
        DELETE FROM api_key_rights
        WHERE api_key_rights_id = $1
        "#,
    )
    .bind(id)
    .execute(pool)
    .await?;
    Ok(result.rows_affected() > 0)
}

pub async fn get_global_api_key_config(pool: &PgPool) -> Result<ApiKeyGlobalConfig, sqlx::Error> {
    let global_config_id = Uuid::nil();

    sqlx::query(
        r#"
        INSERT INTO api_key_config (api_key_config_id, enforce_api_keys)
        VALUES ($1, false)
        ON CONFLICT (api_key_config_id) DO NOTHING
        "#,
    )
    .bind(global_config_id)
    .execute(pool)
    .await?;

    let row: PgRow = sqlx::query(
        r#"
        SELECT enforce_api_keys, updated_at
        FROM api_key_config
        WHERE api_key_config_id = $1
        "#,
    )
    .bind(global_config_id)
    .fetch_one(pool)
    .await?;

    Ok(ApiKeyGlobalConfig {
        enforce_api_keys: row.try_get("enforce_api_keys")?,
        updated_at: row.try_get("updated_at")?,
    })
}

pub async fn set_global_api_key_config(
    pool: &PgPool,
    enforce_api_keys: bool,
) -> Result<ApiKeyGlobalConfig, sqlx::Error> {
    let global_config_id = Uuid::nil();
    sqlx::query(
        r#"
        INSERT INTO api_key_config (api_key_config_id, enforce_api_keys)
        VALUES ($1, $2)
        ON CONFLICT (api_key_config_id)
        DO UPDATE SET
            enforce_api_keys = EXCLUDED.enforce_api_keys,
            updated_at = now()
        "#,
    )
    .bind(global_config_id)
    .bind(enforce_api_keys)
    .execute(pool)
    .await?;

    get_global_api_key_config(pool).await
}

pub async fn list_client_api_key_configs(
    pool: &PgPool,
) -> Result<Vec<ApiKeyClientConfig>, sqlx::Error> {
    let rows: Vec<PgRow> = sqlx::query(
        r#"
        SELECT client_name, enforce_api_keys, created_at, updated_at
        FROM api_key_client_config
        ORDER BY client_name
        "#,
    )
    .fetch_all(pool)
    .await?;

    rows.into_iter()
        .map(|row| {
            Ok(ApiKeyClientConfig {
                client_name: row.try_get("client_name")?,
                enforce_api_keys: row.try_get("enforce_api_keys")?,
                created_at: row.try_get("created_at")?,
                updated_at: row.try_get("updated_at")?,
            })
        })
        .collect()
}

pub async fn get_client_api_key_config(
    pool: &PgPool,
    client_name: &str,
) -> Result<Option<ApiKeyClientConfig>, sqlx::Error> {
    let row: Option<PgRow> = sqlx::query(
        r#"
        SELECT client_name, enforce_api_keys, created_at, updated_at
        FROM api_key_client_config
        WHERE client_name = $1
        "#,
    )
    .bind(client_name)
    .fetch_optional(pool)
    .await?;

    row.map(|row| {
        Ok(ApiKeyClientConfig {
            client_name: row.try_get("client_name")?,
            enforce_api_keys: row.try_get("enforce_api_keys")?,
            created_at: row.try_get("created_at")?,
            updated_at: row.try_get("updated_at")?,
        })
    })
    .transpose()
}

pub async fn upsert_client_api_key_config(
    pool: &PgPool,
    client_name: &str,
    enforce_api_keys: bool,
) -> Result<ApiKeyClientConfig, sqlx::Error> {
    sqlx::query(
        r#"
        INSERT INTO api_key_client_config (api_key_client_config_id, client_name, enforce_api_keys)
        VALUES ($1, $2, $3)
        ON CONFLICT (client_name)
        DO UPDATE SET
            enforce_api_keys = EXCLUDED.enforce_api_keys,
            updated_at = now()
        "#,
    )
    .bind(Uuid::new_v4())
    .bind(client_name)
    .bind(enforce_api_keys)
    .execute(pool)
    .await?;

    get_client_api_key_config(pool, client_name)
        .await?
        .ok_or(sqlx::Error::RowNotFound)
}

pub async fn delete_client_api_key_config(
    pool: &PgPool,
    client_name: &str,
) -> Result<bool, sqlx::Error> {
    let result = sqlx::query(
        r#"
        DELETE FROM api_key_client_config
        WHERE client_name = $1
        "#,
    )
    .bind(client_name)
    .execute(pool)
    .await?;
    Ok(result.rows_affected() > 0)
}

pub async fn get_effective_api_key_enforcement(
    pool: &PgPool,
    client_name: Option<&str>,
) -> Result<bool, sqlx::Error> {
    let global: ApiKeyGlobalConfig = get_global_api_key_config(pool).await?;
    let client_override: Option<ApiKeyClientConfig> =
        if let Some(client_name) = client_name.filter(|value| !value.is_empty()) {
            get_client_api_key_config(pool, client_name).await?
        } else {
            None
        };

    Ok(client_override
        .map(|config| config.enforce_api_keys)
        .unwrap_or(global.enforce_api_keys))
}

pub async fn insert_auth_attempt_log(
    pool: &PgPool,
    entry: AuthAttemptLogEntry,
) -> Result<(), sqlx::Error> {
    let log_id = Uuid::new_v4();
    let column_types = crate::drivers::postgresql::schema_cache::get_public_table_column_types(
        pool,
        "api_key_auth_log",
    )
    .await?;

    let mut query_builder = sqlx::QueryBuilder::<Postgres>::new("INSERT INTO api_key_auth_log (");
    let mut first = true;
    let mut push_column = |query: &mut sqlx::QueryBuilder<'_, Postgres>, column: &str| {
        if !first {
            query.push(", ");
        }
        query.push(column);
        first = false;
    };

    push_column(&mut query_builder, "request_id");
    let has_log_uuid = column_types.contains_key("api_key_auth_log_id");
    if has_log_uuid {
        push_column(&mut query_builder, "api_key_auth_log_id");
    }
    let api_key_id_type = column_types.get("api_key_id").cloned();
    let has_api_key_id = api_key_id_type.is_some();
    if has_api_key_id {
        push_column(&mut query_builder, "api_key_id");
    }
    push_column(&mut query_builder, "api_key_public_id");
    push_column(&mut query_builder, "client_name");
    push_column(&mut query_builder, "method");
    push_column(&mut query_builder, "path");
    push_column(&mut query_builder, "presented_key_hash");
    push_column(&mut query_builder, "presented_key_salt");
    push_column(&mut query_builder, "required_rights");
    push_column(&mut query_builder, "granted_rights");
    push_column(&mut query_builder, "authenticated");
    push_column(&mut query_builder, "authorized");
    push_column(&mut query_builder, "enforced");
    push_column(&mut query_builder, "failure_reason");
    push_column(&mut query_builder, "remote_addr");
    push_column(&mut query_builder, "user_agent");
    let has_time = column_types.contains_key("time");
    if has_time {
        push_column(&mut query_builder, "time");
    }

    query_builder.push(") VALUES (");
    let mut separated = query_builder.separated(", ");
    separated.push_bind(entry.request_id);
    if has_log_uuid {
        separated.push_bind(log_id);
    }
    if let Some(data_type) = api_key_id_type.as_deref() {
        if data_type == "bigint" {
            let internal_id = match entry.internal_api_key_id {
                Some(value) => Some(value),
                None => match entry.api_key_id.as_deref() {
                    Some(external_id) => resolve_internal_api_key_id(pool, external_id).await?,
                    None => None,
                },
            };
            separated.push_bind(internal_id);
        } else {
            separated.push_bind(entry.api_key_id);
        }
    }
    separated.push_bind(entry.api_key_public_id);
    separated.push_bind(entry.client_name);
    separated.push_bind(entry.method);
    separated.push_bind(entry.path);
    separated.push_bind(entry.presented_key_hash);
    separated.push_bind(entry.presented_key_salt);
    separated.push_bind(json!(entry.required_rights));
    separated.push_bind(json!(entry.granted_rights));
    separated.push_bind(entry.authenticated);
    separated.push_bind(entry.authorized);
    separated.push_bind(entry.enforced);
    separated.push_bind(entry.failure_reason);
    separated.push_bind(entry.remote_addr);
    separated.push_bind(entry.user_agent);
    if has_time {
        match column_types.get("time").map(String::as_str) {
            Some("timestamp with time zone") => {
                let timestamp = DateTime::<Utc>::from_timestamp(entry.time, 0);
                separated.push_bind(timestamp);
            }
            _ => {
                separated.push_bind(entry.time);
            }
        };
    }
    separated.push_unseparated(")");

    query_builder.build().execute(pool).await?;
    Ok(())
}

pub async fn resolve_internal_api_key_id(
    pool: &PgPool,
    external_api_key_id: &str,
) -> Result<Option<i64>, sqlx::Error> {
    let Ok(external_api_key_id) = Uuid::parse_str(external_api_key_id) else {
        return Ok(None);
    };

    sqlx::query_scalar(
        r#"
        SELECT id
        FROM api_keys
        WHERE api_keys_id = $1
        "#,
    )
    .bind(external_api_key_id)
    .fetch_optional(pool)
    .await
}

pub fn missing_rights_json(required_rights: &[String], granted_rights: &[String]) -> Value {
    json!({
        "required_rights": required_rights,
        "granted_rights": granted_rights,
    })
}