gitdataai/lib/service/user/sshkey.rs

use base64::{Engine as _, engine::general_purpose};
use db::sqlx;
use model::users::UserSshKeyModel;
use serde::{Deserialize, Serialize};
use session::Session;
use sha2::{Digest, Sha256};

use crate::{AppService, error::AppError, session_user};

#[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)]
pub struct UserSshKey {
    pub id: i64,
    pub title: String,
    pub public_key: String,
    pub fingerprint: String,
    pub key_type: String,
    pub key_bits: Option<i32>,
    pub is_verified: bool,
    #[schema(value_type = Option<String>)]
    pub last_used_at: Option<chrono::DateTime<chrono::Utc>>,
    #[schema(value_type = Option<String>)]
    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
    pub is_revoked: bool,
    #[schema(value_type = String)]
    pub created_at: chrono::DateTime<chrono::Utc>,
    #[schema(value_type = String)]
    pub updated_at: chrono::DateTime<chrono::Utc>,
}

#[derive(Debug, Clone, Deserialize, utoipa::ToSchema)]
pub struct CreateUserSshKey {
    pub title: String,
    pub public_key: String,
    #[schema(value_type = Option<String>)]
    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
}

#[derive(Debug, Clone, Deserialize, utoipa::ToSchema)]
pub struct UpdateUserSshKey {
    pub title: Option<String>,
    #[schema(value_type = Option<String>)]
    pub expires_at: Option<Option<chrono::DateTime<chrono::Utc>>>,
}

impl AppService {
    pub async fn user_ssh_keys(
        &self,
        ctx: &Session,
    ) -> Result<Vec<UserSshKey>, AppError> {
        let user_uid = session_user(ctx)?;
        let rows = sqlx::query_as::<_, UserSshKeyModel>(
            "SELECT id, \"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at \
             FROM user_ssh_key WHERE \"user\" = $1 ORDER BY created_at DESC",
        )
        .bind(user_uid)
        .fetch_all(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        Ok(rows.into_iter().map(Into::into).collect())
    }

    pub async fn user_add_ssh_key(
        &self,
        ctx: &Session,
        params: CreateUserSshKey,
    ) -> Result<UserSshKey, AppError> {
        let user_uid = session_user(ctx)?;
        let title = params.title.trim();
        if title.is_empty() {
            return Err(AppError::BadRequest(
                "ssh key title is required".to_string(),
            ));
        }

        let parsed = parse_public_key(&params.public_key)?;
        let now = chrono::Utc::now();
        let row = sqlx::query_as::<_, UserSshKeyModel>(
            "INSERT INTO user_ssh_key \
                (\"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at) \
             VALUES ($1, $2, $3, $4, $5, $6, true, NULL, $7, false, $8, $8) \
             RETURNING id, \"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at",
        )
        .bind(user_uid)
        .bind(title)
        .bind(parsed.public_key)
        .bind(parsed.fingerprint)
        .bind(parsed.key_type)
        .bind(parsed.key_bits)
        .bind(params.expires_at)
        .bind(now)
        .fetch_one(self.db.writer())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        Ok(row.into())
    }

    pub async fn user_revoke_ssh_key(
        &self,
        ctx: &Session,
        key_id: i64,
    ) -> Result<(), AppError> {
        let user_uid = session_user(ctx)?;
        let affected = sqlx::query(
            "UPDATE user_ssh_key SET is_revoked = true, updated_at = $1 \
             WHERE id = $2 AND \"user\" = $3 AND is_revoked = false",
        )
        .bind(chrono::Utc::now())
        .bind(key_id)
        .bind(user_uid)
        .execute(self.db.writer())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .rows_affected();

        if affected == 0 {
            return Err(AppError::NotFound("ssh key not found".to_string()));
        }
        Ok(())
    }

    pub async fn user_update_ssh_key(
        &self,
        ctx: &Session,
        key_id: i64,
        params: UpdateUserSshKey,
    ) -> Result<UserSshKey, AppError> {
        let user_uid = session_user(ctx)?;
        let mut key = self.user_ssh_key_by_id(user_uid, key_id).await?;

        if let Some(title) = params.title {
            let title = title.trim();
            if title.is_empty() {
                return Err(AppError::BadRequest(
                    "ssh key title is required".to_string(),
                ));
            }
            key.title = title.to_string();
        }
        if let Some(expires_at) = params.expires_at {
            key.expires_at = expires_at;
        }

        let row = sqlx::query_as::<_, UserSshKeyModel>(
            "UPDATE user_ssh_key SET title = $1, expires_at = $2, updated_at = $3 \
             WHERE id = $4 AND \"user\" = $5 AND is_revoked = false \
             RETURNING id, \"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at",
        )
        .bind(&key.title)
        .bind(key.expires_at)
        .bind(chrono::Utc::now())
        .bind(key_id)
        .bind(user_uid)
        .fetch_optional(self.db.writer())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .ok_or_else(|| AppError::NotFound("ssh key not found".to_string()))?;

        Ok(row.into())
    }

    async fn user_ssh_key_by_id(
        &self,
        user_uid: uuid::Uuid,
        key_id: i64,
    ) -> Result<UserSshKey, AppError> {
        sqlx::query_as::<_, UserSshKeyModel>(
            "SELECT id, \"user\", title, public_key, fingerprint, key_type, key_bits, is_verified, last_used_at, expires_at, is_revoked, created_at, updated_at \
             FROM user_ssh_key WHERE id = $1 AND \"user\" = $2 AND is_revoked = false",
        )
        .bind(key_id)
        .bind(user_uid)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .map(Into::into)
        .ok_or_else(|| AppError::NotFound("ssh key not found".to_string()))
    }
}

struct ParsedPublicKey {
    public_key: String,
    fingerprint: String,
    key_type: String,
    key_bits: Option<i32>,
}

fn parse_public_key(public_key: &str) -> Result<ParsedPublicKey, AppError> {
    let public_key = public_key.trim();
    let mut parts = public_key.split_whitespace();
    let key_type = parts.next().ok_or_else(|| {
        AppError::BadRequest("invalid ssh public key".to_string())
    })?;
    let key_data_base64 = parts.next().ok_or_else(|| {
        AppError::BadRequest("invalid ssh public key".to_string())
    })?;

    let key_data =
        general_purpose::STANDARD
            .decode(key_data_base64)
            .map_err(|_| {
                AppError::BadRequest("invalid ssh public key".to_string())
            })?;

    let mut hasher = Sha256::new();
    hasher.update(&key_data);
    let fingerprint = format!(
        "SHA256:{}",
        general_purpose::STANDARD_NO_PAD.encode(hasher.finalize())
    );

    Ok(ParsedPublicKey {
        public_key: public_key.to_string(),
        fingerprint,
        key_type: key_type.to_string(),
        key_bits: key_bits(key_type),
    })
}

fn key_bits(key_type: &str) -> Option<i32> {
    match key_type {
        "ssh-ed25519" => Some(256),
        "ecdsa-sha2-nistp256" => Some(256),
        "ecdsa-sha2-nistp384" => Some(384),
        "ecdsa-sha2-nistp521" => Some(521),
        _ => None,
    }
}

impl From<UserSshKeyModel> for UserSshKey {
    fn from(value: UserSshKeyModel) -> Self {
        Self {
            id: value.id,
            title: value.title,
            public_key: value.public_key,
            fingerprint: value.fingerprint,
            key_type: value.key_type,
            key_bits: value.key_bits,
            is_verified: value.is_verified,
            last_used_at: value.last_used_at,
            expires_at: value.expires_at,
            is_revoked: value.is_revoked,
            created_at: value.created_at,
            updated_at: value.updated_at,
        }
    }
}