dynamic-waas-sdk 0.0.3

//! Encrypted-share backup orchestration.
//!
//! Mirrors `python/dynamic_wallet_sdk/wallet_client.py::backup_key_shares`:
//!
//! 1. Encrypt each `ServerKeyShare`'s `secret_share` with the customer-
//!    provided password using AES-256-GCM (PBKDF2-SHA256 v2 params from
//!    `crate::crypto`). The plaintext is a JSON wrapper carrying both
//!    `keyShareId` and `secretShare` so recovery can round-trip them.
//! 2. Base64-wrap the JSON-serialised `EncryptedData` (`salt/iv/cipher/
//!    version`) and POST the list as `encryptedAccountCredentials` to
//!    `/waas/{walletId}/keyShares/backup`.
//! 3. The server hands back `keyShareIds` — one server-assigned UUID per
//!    encrypted blob, in the same order.
//! 4. POST `/waas/{walletId}/keyShares/backup/locations` with one
//!    `BackupLocation::Dynamic` entry per share, using the server-assigned
//!    `externalKeyShareId` and the MPC `keygen_id` of the first share.
//!    This is the step that flips the wallet from "pending" to "active";
//!    until it runs, the relay rejects signing.
//! 5. Return a `KeyShareBackupInfo` carrying the `password_encrypted`
//!    flag and the per-location pointers so the caller can merge it into
//!    their cached `WalletProperties`.
//!
//! Encryption-version `v2` (1M PBKDF2 iterations) is hard-coded because
//! Dynamic's recovery endpoint reads the version from the blob itself —
//! we send `encryptionVersion=v2` in the metadata so the relay can route
//! pre-decrypt heuristics correctly. v1 (100k iterations) only exists
//! for legacy backups; new backups always use v2.

use base64::engine::general_purpose::STANDARD as B64;
use base64::Engine;
use dynamic_waas_sdk_core::{
    api::{RecoverBackupReq, StoreBackupReq},
    BackupLocation, BackupLocationInfo, Error, KeyShareBackupInfo, Result, ServerKeyShare,
};

use crate::client::DynamicWalletClient;
use crate::crypto::{self, EncryptedData};

/// JSON wrapper for the plaintext that goes into the encrypted blob.
#[derive(serde::Serialize, serde::Deserialize)]
struct KeyShareWrapper<'a> {
    #[serde(rename = "keyShareId")]
    key_share_id: &'a str,
    #[serde(rename = "secretShare")]
    secret_share: &'a str,
}

/// JSON shape `EncryptedData` deserialises to/from on the wire.
#[derive(Debug, serde::Deserialize)]
struct RecoveredKeyShareWrapper {
    #[serde(rename = "keyShareId")]
    key_share_id: Option<String>,
    #[serde(rename = "secretShare")]
    secret_share: String,
}

fn encrypt_share(share: &ServerKeyShare, password: &str) -> Result<String> {
    let wrapper = KeyShareWrapper {
        key_share_id: &share.key_share_id,
        secret_share: &share.secret_share,
    };
    let plaintext = serde_json::to_string(&wrapper)?;
    let enc = crypto::encrypt(&plaintext, password)?;
    let enc_json = serde_json::to_string(&enc)?;
    Ok(B64.encode(enc_json.as_bytes()))
}

fn decode_recovered_blob(blob: &str, password: &str) -> Result<RecoveredKeyShareWrapper> {
    // Server stores the credential base64-encoded; same shape on the way back.
    let enc_json_bytes = B64
        .decode(blob)
        .map_err(|e| Error::Encryption(format!("backup blob b64 decode: {e}")))?;
    let enc_json = std::str::from_utf8(&enc_json_bytes)
        .map_err(|e| Error::Encryption(format!("backup blob not UTF-8: {e}")))?;
    let enc: EncryptedData = serde_json::from_str(enc_json)?;
    let plaintext = crypto::decrypt(&enc, password)?;
    serde_json::from_str::<RecoveredKeyShareWrapper>(&plaintext).map_err(Error::from)
}

/// Encrypt the customer's shares with `password`, upload them to Dynamic's
/// backup store, then mark the wallet active with one
/// `BackupLocation::Dynamic` per share.
///
/// `keygen_id` is the MPC `export_id` of the first share — chain clients
/// compute it via their respective signer (`EcdsaSigner::export_id` for
/// EVM, `Ed25519Signer::export_id` for SVM) before calling this fn.
///
/// Returns the new `KeyShareBackupInfo` the caller merges into their
/// cached `WalletProperties`.
pub async fn run_backup_dynamic(
    client: &DynamicWalletClient,
    wallet_id: &str,
    shares: &[ServerKeyShare],
    keygen_id: &str,
    password: &str,
) -> Result<KeyShareBackupInfo> {
    if shares.is_empty() {
        return Err(Error::InvalidArgument(
            "shares is empty — nothing to back up".into(),
        ));
    }

    let encrypted_blobs: Result<Vec<String>> =
        shares.iter().map(|s| encrypt_share(s, password)).collect();
    let encrypted_account_credentials = encrypted_blobs?;

    let store_resp = client
        .api()
        .store_encrypted_backup(
            wallet_id,
            &StoreBackupReq {
                encrypted_account_credentials,
                password_encrypted: true,
                encryption_version: Some(crypto::V2.to_string()),
            },
        )
        .await?;
    let server_key_share_ids = store_resp.key_share_ids;

    // Build locations with the SERVER-assigned ids (not our client keygen
    // ids). The relay validates that every share-slot has a corresponding
    // externalKeyShareId from this exact response.
    let locations: Vec<serde_json::Value> = server_key_share_ids
        .iter()
        .map(|id| {
            serde_json::json!({
                "location": BackupLocation::Dynamic.as_wire_str(),
                "passwordEncrypted": true,
                "keygenId": keygen_id,
                "externalKeyShareId": id,
            })
        })
        .collect();
    let mark_body = serde_json::json!({ "locations": locations });
    client
        .api()
        .mark_key_shares_as_backed_up(wallet_id, &mark_body)
        .await?;

    let backups = std::iter::once((
        BackupLocation::Dynamic.as_wire_str().to_string(),
        server_key_share_ids
            .into_iter()
            .map(|id| BackupLocationInfo {
                location: BackupLocation::Dynamic,
                key_share_id: id,
            })
            .collect::<Vec<_>>(),
    ))
    .collect();

    Ok(KeyShareBackupInfo {
        password_encrypted: true,
        backups,
    })
}

/// Mark the wallet active without uploading anything to Dynamic's backup
/// store. Used when the caller has set `back_up_to_dynamic=false` —
/// they're vaulting shares themselves and only need the relay to flip
/// the wallet from "pending" → "active".
///
/// `keygen_id` is the MPC `export_id` of the first share, same as
/// [`run_backup_dynamic`].
pub async fn run_mark_external_no_backup(
    client: &DynamicWalletClient,
    wallet_id: &str,
    shares: &[ServerKeyShare],
    keygen_id: &str,
) -> Result<()> {
    if shares.is_empty() {
        return Err(Error::InvalidArgument(
            "shares is empty — nothing to mark".into(),
        ));
    }
    let locations: Vec<serde_json::Value> = shares
        .iter()
        .map(|_| {
            serde_json::json!({
                "location": BackupLocation::External.as_wire_str(),
                "passwordEncrypted": false,
                "keygenId": keygen_id,
            })
        })
        .collect();
    let mark_body = serde_json::json!({ "locations": locations });
    client
        .api()
        .mark_key_shares_as_backed_up(wallet_id, &mark_body)
        .await?;
    Ok(())
}

/// Recover the customer's shares from Dynamic's backup store, decrypt
/// with `password`, and return them. The caller re-vaults the result.
///
/// `key_share_ids` is the list of server-assigned UUIDs from the original
/// backup (typically read from
/// `wallet_properties.external_server_key_shares_backup_info.backups
/// ["dynamic"]`).
pub async fn run_recover_key_shares(
    client: &DynamicWalletClient,
    wallet_id: &str,
    key_share_ids: Vec<String>,
    password: &str,
) -> Result<Vec<ServerKeyShare>> {
    let resp = client
        .api()
        .recover_encrypted_backup(wallet_id, &RecoverBackupReq { key_share_ids })
        .await?;

    let mut recovered = Vec::with_capacity(resp.key_shares.len());
    for entry in &resp.key_shares {
        let Some(blob) = entry.encrypted_blob() else {
            continue;
        };
        let wrapper = decode_recovered_blob(blob, password)?;
        let key_share_id = wrapper
            .key_share_id
            .or_else(|| entry.id.clone())
            .unwrap_or_default();
        recovered.push(ServerKeyShare::new(key_share_id, wrapper.secret_share));
    }

    if recovered.is_empty() {
        return Err(Error::InvalidArgument(
            "Key share recovery returned no shares. Check that the password is \
             correct and the backup exists."
                .into(),
        ));
    }

    Ok(recovered)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn encrypt_share_round_trips_through_decode() {
        let share = ServerKeyShare::new("ks-1", "secret-data");
        let blob = encrypt_share(&share, "pw").unwrap();

        let wrapper = decode_recovered_blob(&blob, "pw").unwrap();
        assert_eq!(wrapper.key_share_id.as_deref(), Some("ks-1"));
        assert_eq!(wrapper.secret_share, "secret-data");
    }

    #[test]
    fn decode_rejects_wrong_password() {
        let share = ServerKeyShare::new("ks-1", "secret-data");
        let blob = encrypt_share(&share, "right-pw").unwrap();

        let err = decode_recovered_blob(&blob, "wrong-pw").unwrap_err();
        assert!(matches!(err, Error::Encryption(_)));
    }
}