Add More Metadata to public ssh keys (#1182)

Co-authored-by: Eugene <inbox@null.page>

warpgate-admin/src/api/public_key_credentials.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use chrono::{DateTime, Utc};
 use poem::web::Data;
 use poem_openapi::param::Path;
 use poem_openapi::payload::Json;
@@ -19,6 +20,8 @@ use super::AnySecurityScheme;
 struct ExistingPublicKeyCredential {
     id: Uuid,
     label: String,
+    date_added: Option<DateTime<Utc>>,
+    last_used: Option<DateTime<Utc>>,
     openssh_public_key: String,
 }
 
@@ -32,6 +35,8 @@ impl From<PublicKeyCredential::Model> for ExistingPublicKeyCredential {
     fn from(credential: PublicKeyCredential::Model) -> Self {
         Self {
             id: credential.id,
+            date_added: credential.date_added,
+            last_used: credential.last_used,
             label: credential.label,
             openssh_public_key: credential.openssh_public_key,
         }
@@ -115,6 +120,8 @@ impl ListApi {
         let object = PublicKeyCredential::ActiveModel {
             id: Set(Uuid::new_v4()),
             user_id: Set(*user_id),
+            date_added: Set(Some(Utc::now())),
+            last_used: Set(None),
             label: Set(body.label.clone()),
             ..PublicKeyCredential::ActiveModel::from(UserPublicKeyCredential::try_from(&*body)?)
         }
@@ -158,6 +165,7 @@ impl DetailApi {
         let model = PublicKeyCredential::ActiveModel {
             id: Set(id.0),
             user_id: Set(*user_id),
+            date_added: Set(Some(Utc::now())),
             label: Set(body.label.clone()),
             ..<_>::from(UserPublicKeyCredential::try_from(&*body)?)
         }

warpgate-core/src/config_providers/db.rs

@@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use chrono::Utc;
 use data_encoding::BASE64;
 use sea_orm::{
     ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, ModelTrait, QueryFilter,
@@ -381,4 +382,58 @@ impl ConfigProvider for DatabaseConfigProvider {
 
         Ok(())
     }
+
+    async fn update_public_key_last_used(
+        &self,
+        credential: Option<AuthCredential>,
+    ) -> Result<(), WarpgateError> {
+        let db = self.db.lock().await;
+
+        let Some(AuthCredential::PublicKey {
+            kind,
+            public_key_bytes,
+        }) = credential
+        else {
+            error!("Invalid or missing public key credential");
+            return Err(WarpgateError::InvalidCredentialType);
+        };
+
+        // Encode public key and match it against the database
+        let base64_bytes = data_encoding::BASE64.encode(&public_key_bytes);
+        let openssh_public_key = format!("{kind} {base64_bytes}");
+
+        debug!(
+            "Attempting to update last_used for public key: {}",
+            openssh_public_key
+        );
+
+        // Find the public key credential
+        let public_key_credential = entities::PublicKeyCredential::Entity::find()
+            .filter(
+                entities::PublicKeyCredential::Column::OpensshPublicKey
+                    .eq(openssh_public_key.clone()),
+            )
+            .one(&*db)
+            .await?;
+
+        let Some(public_key_credential) = public_key_credential else {
+            warn!(
+                "Public key not found in the database: {}",
+                openssh_public_key
+            );
+            return Ok(()); // Gracefully return if the key is not found
+        };
+
+        // Update the `last_used` (last used) timestamp
+        let mut active_model: entities::PublicKeyCredential::ActiveModel =
+            public_key_credential.into();
+        active_model.last_used = Set(Some(Utc::now()));
+
+        active_model.update(&*db).await.map_err(|e| {
+            error!("Failed to update last_used for public key: {:?}", e);
+            WarpgateError::DatabaseError(e.into())
+        })?;
+
+        Ok(())
+    }
 }

warpgate-core/src/config_providers/mod.rs

@@ -47,6 +47,11 @@ pub trait ConfigProvider {
         username: &str,
         target: &str,
     ) -> Result<bool, WarpgateError>;
+
+    async fn update_public_key_last_used(
+        &self,
+        credential: Option<AuthCredential>,
+    ) -> Result<(), WarpgateError>;
 }
 
 //TODO: move this somewhere

warpgate-db-entities/src/PublicKeyCredential.rs

@@ -1,3 +1,4 @@
+use chrono::{DateTime, Utc};
 use sea_orm::entity::prelude::*;
 use sea_orm::sea_query::ForeignKeyAction;
 use sea_orm::Set;
@@ -12,6 +13,8 @@ pub struct Model {
     pub id: Uuid,
     pub user_id: Uuid,
     pub label: String,
+    pub date_added: Option<DateTime<Utc>>,
+    pub last_used: Option<DateTime<Utc>>,
     pub openssh_public_key: String,
 }
 

warpgate-db-migrations/src/lib.rs

@@ -14,6 +14,7 @@ mod m00009_credential_models;
 mod m00010_parameters;
 mod m00011_rsa_key_algos;
 mod m00012_add_openssh_public_key_label;
+mod m00013_add_openssh_public_key_dates;
 
 pub struct Migrator;
 
@@ -33,6 +34,7 @@ impl MigratorTrait for Migrator {
             Box::new(m00010_parameters::Migration),
             Box::new(m00011_rsa_key_algos::Migration),
             Box::new(m00012_add_openssh_public_key_label::Migration),
+            Box::new(m00013_add_openssh_public_key_dates::Migration),
         ]
     }
 }

warpgate-db-migrations/src/m00012_add_openssh_public_key_label.rs

@@ -21,9 +21,9 @@ impl MigrationTrait for Migration {
                         ColumnDef::new(Alias::new("label"))
                             .string()
                             .not_null()
-                            .default("Public Key")
+                            .default("Public Key"),
                     )
-                    .to_owned()
+                    .to_owned(),
             )
             .await
     }
@@ -38,5 +38,4 @@ impl MigrationTrait for Migration {
             )
             .await
     }
-
 }

warpgate-db-migrations/src/m00013_add_openssh_public_key_dates.rs

@@ -0,0 +1,62 @@
+use sea_orm_migration::prelude::*;
+
+pub struct Migration;
+
+impl MigrationName for Migration {
+    fn name(&self) -> &str {
+        "m00013_add_openssh_public_key_dates"
+    }
+}
+
+use crate::m00009_credential_models::public_key_credential;
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        // Add 'date_added' column
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(public_key_credential::Entity)
+                    .add_column(ColumnDef::new(Alias::new("date_added")).date_time().null())
+                    .to_owned(),
+            )
+            .await?;
+
+        // Add 'last_used' column
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(public_key_credential::Entity)
+                    .add_column(ColumnDef::new(Alias::new("last_used")).date_time().null())
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        // Drop 'last_used' column
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(public_key_credential::Entity)
+                    .drop_column(Alias::new("last_used"))
+                    .to_owned(),
+            )
+            .await?;
+
+        // Drop 'date_added' column
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(public_key_credential::Entity)
+                    .drop_column(Alias::new("date_added"))
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+}

warpgate-protocol-http/src/api/credentials.rs

@@ -1,3 +1,4 @@
+use chrono::{DateTime, Utc};
 use http::StatusCode;
 use poem::web::Data;
 use poem::{Endpoint, EndpointExt, FromRequest, IntoResponse};
@@ -80,6 +81,8 @@ struct NewPublicKeyCredential {
 struct ExistingPublicKeyCredential {
     id: Uuid,
     label: String,
+    date_added: Option<DateTime<Utc>>,
+    last_used: Option<DateTime<Utc>>,
     abbreviated: String,
 }
 
@@ -91,8 +94,8 @@ fn abbreviate_public_key(k: &str) -> String {
 
     format!(
         "{}...{}",
-        &k[..l.min(k.len())],                // Take the first `l` characters.
-        &k[k.len().saturating_sub(l)..]      // Take the last `l` characters safely.
+        &k[..l.min(k.len())],            // Take the first `l` characters.
+        &k[k.len().saturating_sub(l)..]  // Take the last `l` characters safely.
     )
 }
 
@@ -101,6 +104,8 @@ impl From<entities::PublicKeyCredential::Model> for ExistingPublicKeyCredential
         Self {
             id: credential.id,
             label: credential.label,
+            date_added: credential.date_added,
+            last_used: credential.last_used,
             abbreviated: abbreviate_public_key(&credential.openssh_public_key),
         }
     }
@@ -295,6 +300,8 @@ impl Api {
         let object = PublicKeyCredential::ActiveModel {
             id: Set(Uuid::new_v4()),
             user_id: Set(user_model.id),
+            date_added: Set(Some(Utc::now())),
+            last_used: Set(None),
             label: Set(body.label.clone()),
             openssh_public_key: Set(body.openssh_public_key.clone()),
         }

warpgate-protocol-ssh/src/server/session.rs

@@ -1222,18 +1222,28 @@ impl ServerSession {
             key.public_key_base64()
         );
 
-        let result = self
-            .try_auth_lazy(
-                &selector,
-                Some(AuthCredential::PublicKey {
-                    kind: key.algorithm(),
-                    public_key_bytes: Bytes::from(key.public_key_bytes()),
-                }),
-            )
-            .await;
+        let key = Some(AuthCredential::PublicKey {
+            kind: key.algorithm(),
+            public_key_bytes: Bytes::from(key.public_key_bytes()),
+        });
+
+        let result = self.try_auth_lazy(&selector, key.clone()).await;
 
         match result {
-            Ok(AuthResult::Accepted { .. }) => russh::server::Auth::Accept,
+            Ok(AuthResult::Accepted { .. }) => {
+                // Update last_used timestamp
+                if let Err(err) = self
+                    .services
+                    .config_provider
+                    .lock()
+                    .await
+                    .update_public_key_last_used(key.clone())
+                    .await
+                {
+                    warn!(?err, "Failed to update last_used for public key");
+                }
+                russh::server::Auth::Accept
+            }
             Ok(AuthResult::Rejected) => russh::server::Auth::Reject {
                 proceed_with_methods: Some(MethodSet::all()),
             },

warpgate-web/src/admin/CredentialEditor.svelte

@@ -20,6 +20,7 @@
     import CreateOtpModal from './CreateOtpModal.svelte'
     import AuthPolicyEditor from './AuthPolicyEditor.svelte'
     import { possibleCredentials } from 'common/protocols'
+    import CredentialUsedStateBadge from 'common/CredentialUsedStateBadge.svelte'
 
     interface Props {
         userId: string
@@ -248,27 +249,32 @@
     <div class="list-group-item credential">
         {#if credential.kind === CredentialKind.Password }
             <Fa fw icon={faKeyboard} />
-            <span class="type">Password</span>
+            <span class="label me-auto">Password</span>
         {/if}
         {#if credential.kind === 'PublicKey'}
             <Fa fw icon={faKey} />
-            <span class="type">{credential.label}</span>
-            <span class="text-muted ms-2">{abbreviatePublicKey(credential.opensshPublicKey)}</span>
+            <div class="main me-auto">
+                <div class="label d-flex align-items-center">
+                    {credential.label}
+                </div>
+                <small class="d-block text-muted">{abbreviatePublicKey(credential.opensshPublicKey)}</small>
+            </div>
+            <CredentialUsedStateBadge credential={credential} />
+            <div class="me-2"></div>
         {/if}
         {#if credential.kind === 'Totp'}
             <Fa fw icon={faMobileScreen} />
-            <span class="type">One-time password</span>
+            <span class="label me-auto">One-time password</span>
         {/if}
         {#if credential.kind === CredentialKind.Sso}
             <Fa fw icon={faIdBadge} />
-            <span class="type">Single sign-on</span>
-            <span class="text-muted ms-2">
+            <span class="label">Single sign-on</span>
+            <span class="text-muted ms-2 me-auto">
                 {credential.email}
                 {#if credential.provider} ({credential.provider}){/if}
             </span>
         {/if}
 
-        <span class="ms-auto"></span>
         {#if credential.kind === CredentialKind.PublicKey || credential.kind === CredentialKind.Sso}
         <a
             class="ms-2"
@@ -362,16 +368,8 @@
         display: flex;
         align-items: center;
 
-        .type {
-            margin-left: .5rem;
-        }
-
-        a {
-            display: none;
-        }
-
-        &:hover a {
-            display: initial;
+        .label:not(:first-child), .main {
+            margin-left: .75rem;
         }
     }
 </style>

warpgate-web/src/admin/Session.svelte

@@ -82,8 +82,8 @@
         <div>
             <h1>session</h1>
             <div class="d-flex align-items-center mt-1">
-                <Tooltip delay="500" target="usernameBadge" animation>Authenticated user</Tooltip>
-                <Tooltip delay="500" target="targetBadge" animation>Selected target</Tooltip>
+                <Tooltip delay="250" target="usernameBadge" animation>Authenticated user</Tooltip>
+                <Tooltip delay="250" target="targetBadge" animation>Selected target</Tooltip>
 
                 <Badge id="usernameBadge" color="success" class="me-2 d-flex align-items-center">
                     {#if session.username}