kpx 0.0.1

#[allow(deprecated)]
use crypto_box::aead::generic_array::GenericArray;
use crypto_box::aead::{Aead, AeadCore};
use crypto_box::{PublicKey, SalsaBox, SecretKey};
use rand_core::OsRng;
use serde_json::{Map, Value, json};
use tracing::{instrument, trace};

use crate::{
    error::KpxError,
    models::{
        Action, Association, AssociationRecord, DatabaseGroups, DatabaseHash, Envelope,
        GetLoginsResponse, LoginQuery, PasswordResponse, SuccessValue, TestAssociationResult,
        TotpResponse,
    },
    transport::{NativeTransport, Transport},
    util::{
        CLIENT_ID_LEN, KEY_LEN, NONCE_LEN, decode_base64, decode_base64_array, encode_base64,
        random_bytes, success_to_bool,
    },
};

/// Convenience alias for the default KeePassXC client leveraging the native transport.
pub type KpxClient = KeePassXcClient<NativeTransport>;

/// Synchronous KeePassXC protocol client.
pub struct KeePassXcClient<T: Transport> {
    transport: T,
    secret_key: SecretKey,
    public_key: PublicKey,
    host_public_key: Option<PublicKey>,
    client_id_b64: String,
    associations: Vec<AssociationRecord>,
}

impl KeePassXcClient<NativeTransport> {
    /// Connect to KeePassXC using the platform native transport.
    #[instrument(level = "info", err)]
    pub fn connect() -> Result<Self, KpxError> {
        let transport = NativeTransport::connect_default()?;
        Ok(Self::new(transport))
    }
}

impl<T: Transport> KeePassXcClient<T> {
    #[instrument(level = "debug", skip(transport))]
    /// Create a KeePassXC client using the provided transport implementation.
    pub fn new(transport: T) -> Self {
        let (public_key, secret_key) = generate_keypair();
        let client_id_b64 = {
            let bytes = random_bytes::<CLIENT_ID_LEN>();
            encode_base64(&bytes)
        };

        Self {
            transport,
            secret_key,
            public_key,
            host_public_key: None,
            client_id_b64,
            associations: Vec::new(),
        }
    }

    /// Return the base64-encoded identifier negotiated with KeePassXC.
    pub fn client_id(&self) -> &str {
        &self.client_id_b64
    }

    /// Return all association records currently cached by the client.
    pub fn associations(&self) -> &[AssociationRecord] {
        &self.associations
    }

    #[instrument(level = "debug", skip(self), err)]
    /// Perform the public key exchange with KeePassXC and cache the host key.
    pub fn change_public_keys(&mut self) -> Result<Option<String>, KpxError> {
        let nonce = random_bytes::<NONCE_LEN>();
        let request = RequestBuilder::new(Action::ChangePublicKeys)
            .with_client_id(&self.client_id_b64)
            .with_nonce(&nonce)
            .with_value(
                "publicKey",
                Value::String(encode_base64(self.public_key.as_bytes())),
            )
            .build();

        let envelope = self.send_plain(Value::Object(request))?;
        ensure_action(Action::ChangePublicKeys, &envelope)?;

        let public_key_b64 = envelope
            .public_key
            .as_deref()
            .ok_or(KpxError::MissingField("publicKey"))?;
        let host_key_bytes = decode_base64_array::<KEY_LEN>(public_key_b64)?;
        let host_public_key = PublicKey::from(host_key_bytes);
        self.host_public_key = Some(host_public_key);

        let success = success_to_bool(envelope.success.as_ref());
        if matches!(success, Some(false)) {
            return Err(KpxError::RemoteError {
                code: envelope.error_code.as_ref().and_then(|c| c.as_i32()),
                message: envelope
                    .error
                    .clone()
                    .unwrap_or_else(|| "change-public-keys rejected".to_string()),
            });
        }

        Ok(envelope.version.clone())
    }

    #[instrument(level = "info", skip(self), err)]
    /// Create a new association in KeePassXC and persist it in memory.
    pub fn associate(&mut self) -> Result<AssociationRecord, KpxError> {
        self.ensure_host_key()?;

        let (id_public, _) = generate_keypair();
        let id_key_b64 = encode_base64(id_public.as_bytes());
        let session_key_b64 = encode_base64(self.public_key.as_bytes());

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::Associate.as_str().to_string()),
        );
        payload.insert("key".into(), Value::String(session_key_b64.clone()));
        payload.insert("idKey".into(), Value::String(id_key_b64.clone()));

        let payload =
            self.send_encrypted_payload(Action::Associate, Some(Value::Object(payload)), None)?;

        // Check for errors first - if there's an error field, the request failed
        if let Some(error_msg) = payload.get("error").and_then(Value::as_str) {
            return Err(KpxError::AssociationRejected {
                reason: error_msg.to_string(),
            });
        }

        let success = parse_success(&payload);
        if let Some(false) = success {
            let reason = payload
                .get("error")
                .and_then(Value::as_str)
                .unwrap_or("association rejected")
                .to_string();
            return Err(KpxError::AssociationRejected { reason });
        }

        let id = payload
            .get("id")
            .and_then(Value::as_str)
            .ok_or(KpxError::MissingField("id"))?
            .to_string();
        let database_hash = payload
            .get("hash")
            .and_then(Value::as_str)
            .ok_or(KpxError::MissingField("hash"))?
            .to_string();
        let version = payload
            .get("version")
            .and_then(Value::as_str)
            .map(|value| value.to_string())
            .unwrap_or_default();
        let version_str = if version.is_empty() {
            None
        } else {
            Some(version.as_str())
        };

        let association_key = if supports_identification_key(version_str) {
            id_key_b64
        } else {
            session_key_b64
        };

        let record = AssociationRecord {
            association: Association {
                id,
                key: association_key,
            },
            database_hash,
            version,
        };
        self.add_association(record.clone());

        Ok(record)
    }

    #[instrument(
        level = "debug",
        skip(self, association),
        err,
        fields(association_id = %association.id)
    )]
    /// Verify that an existing association is still accepted by KeePassXC.
    pub fn test_associate(
        &mut self,
        association: &Association,
    ) -> Result<TestAssociationResult, KpxError> {
        self.ensure_host_key()?;

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::TestAssociate.as_str().to_string()),
        );
        payload.insert("id".into(), Value::String(association.id.clone()));
        payload.insert("key".into(), Value::String(association.key.clone()));

        let payload =
            self.send_encrypted_payload(Action::TestAssociate, Some(Value::Object(payload)), None)?;
        let result: TestAssociationResult = serde_json::from_value(payload)?;
        Ok(result)
    }

    #[instrument(level = "debug", skip(self, request_id), err)]
    /// Ask KeePassXC to generate a password, optionally scoped by a `request_id`.
    pub fn generate_password(
        &mut self,
        request_id: Option<String>,
    ) -> Result<PasswordResponse, KpxError> {
        self.ensure_host_key()?;

        let mut extra = Map::new();
        if let Some(request_id) = request_id {
            extra.insert("requestID".into(), Value::String(request_id));
        }

        let payload = self.send_encrypted_payload(Action::GeneratePassword, None, Some(extra))?;

        let payload: PasswordResponse = serde_json::from_value(payload)?;
        Ok(payload)
    }

    #[instrument(
        level = "info",
        skip(self, query),
        err,
        fields(url = %query.url, keys = query.keys.len())
    )]
    /// Retrieve logins for the provided query, decrypting the entries transparently.
    pub fn get_logins(&mut self, query: &LoginQuery) -> Result<GetLoginsResponse, KpxError> {
        self.ensure_host_key()?;

        let keys = if query.keys.is_empty() {
            self.associations
                .iter()
                .map(|record| record.association.clone())
                .collect()
        } else {
            query.keys.clone()
        };

        if keys.is_empty() {
            return Err(KpxError::MissingField("keys"));
        }

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::GetLogins.as_str().to_string()),
        );
        payload.insert("url".into(), Value::String(query.url.clone()));
        if let Some(submit) = &query.submit_url {
            payload.insert("submitUrl".into(), Value::String(submit.clone()));
        }
        if let Some(auth) = &query.http_auth {
            payload.insert("httpAuth".into(), Value::String(auth.clone()));
        }
        if let Some(primary_id) = &query.primary_id {
            payload.insert("id".into(), Value::String(primary_id.clone()));
        }

        let keys_json: Vec<Value> = keys
            .into_iter()
            .map(|assoc| json!({ "id": assoc.id, "key": assoc.key }))
            .collect();
        payload.insert("keys".into(), Value::Array(keys_json));

        let payload =
            self.send_encrypted_payload(Action::GetLogins, Some(Value::Object(payload)), None)?;
        let result: GetLoginsResponse = serde_json::from_value(payload)?;
        Ok(result)
    }

    #[instrument(level = "debug", skip(self), err)]
    /// Fetch the KeePassXC database group tree.
    pub fn get_database_groups(&mut self) -> Result<DatabaseGroups, KpxError> {
        self.ensure_host_key()?;

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::GetDatabaseGroups.as_str().to_string()),
        );

        let payload = self.send_encrypted_payload(
            Action::GetDatabaseGroups,
            Some(Value::Object(payload)),
            None,
        )?;
        let groups: DatabaseGroups = serde_json::from_value(payload)?;
        Ok(groups)
    }

    #[instrument(level = "debug", skip(self), err)]
    /// Retrieve the KeePassXC database hash to identify the open database.
    pub fn get_database_hash(&mut self) -> Result<DatabaseHash, KpxError> {
        self.ensure_host_key()?;

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::GetDatabaseHash.as_str().to_string()),
        );

        let payload = self.send_encrypted_payload(
            Action::GetDatabaseHash,
            Some(Value::Object(payload)),
            None,
        )?;
        let hash: DatabaseHash = serde_json::from_value(payload)?;
        Ok(hash)
    }

    #[instrument(level = "debug", skip(self), err, fields(uuid = %uuid))]
    /// Request a TOTP code for the entry identified by `uuid`.
    pub fn get_totp(&mut self, uuid: &str) -> Result<Option<String>, KpxError> {
        self.ensure_host_key()?;

        let mut payload = Map::new();
        payload.insert(
            "action".into(),
            Value::String(Action::GetTotp.as_str().to_string()),
        );
        payload.insert("uuid".into(), Value::String(uuid.to_string()));

        let payload =
            self.send_encrypted_payload(Action::GetTotp, Some(Value::Object(payload)), None)?;
        if payload.is_null() {
            return Ok(None);
        }

        let response: TotpResponse = serde_json::from_value(payload)?;
        if let Some(false) = success_to_bool(response.success.as_ref()) {
            return Err(KpxError::RemoteError {
                code: response.error_code.as_ref().and_then(|c| c.as_i32()),
                message: response
                    .error
                    .unwrap_or_else(|| "get-totp rejected".to_string()),
            });
        }

        Ok(response.totp)
    }

    #[instrument(level = "trace", skip(self, record), fields(id = %record.association.id))]
    /// Cache a new association record if it has not been seen yet.
    pub fn add_association(&mut self, record: AssociationRecord) {
        if !self
            .associations
            .iter()
            .any(|existing| existing.association == record.association)
        {
            self.associations.push(record);
        }
    }

    #[instrument(level = "trace", skip(self), err)]
    fn ensure_host_key(&mut self) -> Result<(), KpxError> {
        if self.host_public_key.is_none() {
            self.change_public_keys()?;
        }
        Ok(())
    }

    #[instrument(level = "trace", skip(self, request), err)]
    fn send_plain(&mut self, request: Value) -> Result<Envelope, KpxError> {
        let serialized = serde_json::to_string(&request)?;
        let action_label = request
            .get("action")
            .and_then(|value| value.as_str())
            .unwrap_or("unknown");
        trace!(action = action_label, request = %serialized, "sending KeePassXC request");
        self.transport.send_line(&serialized)?;
        let line = self.transport.read_line()?;
        trace!(response = %line, "received KeePassXC response");
        if line.trim().is_empty() {
            return Err(KpxError::EmptyResponse);
        }
        let envelope: Envelope = serde_json::from_str(&line)?;
        Ok(envelope)
    }

    #[instrument(
        level = "trace",
        skip(self, payload, extra_fields),
        err,
        fields(action = %action.as_str())
    )]
    fn send_encrypted_payload(
        &mut self,
        action: Action,
        payload: Option<Value>,
        extra_fields: Option<Map<String, Value>>,
    ) -> Result<Value, KpxError> {
        self.ensure_host_key()?;
        let host_public = self
            .host_public_key
            .as_ref()
            .ok_or(KpxError::MissingHostKey)?
            .clone();
        let salsabox = SalsaBox::new(&host_public, &self.secret_key);

        let nonce_bytes = random_bytes::<NONCE_LEN>();
        let mut request_builder = RequestBuilder::new(action)
            .with_client_id(&self.client_id_b64)
            .with_nonce(&nonce_bytes);

        if let Some(extra) = extra_fields {
            request_builder = request_builder.extend(extra);
        }

        if let Some(payload) = payload {
            let payload_bytes = serde_json::to_vec(&payload)?;
            let nonce = make_nonce(&nonce_bytes);
            let ciphertext = salsabox.encrypt(&nonce, payload_bytes.as_slice())?;
            request_builder =
                request_builder.with_value("message", Value::String(encode_base64(&ciphertext)));
            trace!(
                action = %action.as_str(),
                payload = %serde_json::to_string(&payload).unwrap_or_else(|_| "<unserializable>".to_string()),
                "encrypted KeePassXC payload"
            );
        }

        let envelope = self.send_plain(Value::Object(request_builder.build()))?;
        ensure_action(action, &envelope)?;

        // Check if the envelope contains an error before trying to decrypt
        if let Some(error_msg) = &envelope.error {
            return Err(KpxError::RemoteError {
                code: envelope.error_code.as_ref().and_then(|c| c.as_i32()),
                message: error_msg.clone(),
            });
        }

        let message_b64 = match envelope.message.clone() {
            Some(message) => message,
            None => return Ok(Value::Null),
        };

        let nonce_b64 = envelope
            .nonce
            .as_deref()
            .ok_or(KpxError::MissingField("nonce"))?;
        let nonce_bytes = decode_base64_array::<NONCE_LEN>(nonce_b64)?;
        let ciphertext = decode_base64(&message_b64)?;

        let nonce = make_nonce(&nonce_bytes);
        let plaintext = salsabox.decrypt(&nonce, ciphertext.as_slice())?;
        let payload_json: Value = serde_json::from_slice(&plaintext)?;
        trace!(action = %action.as_str(), payload = %payload_json, "decrypted KeePassXC payload");

        Ok(payload_json)
    }
}

#[instrument(level = "trace", skip(envelope), err, fields(expected = %expected.as_str(), received = %envelope.action))]
fn ensure_action(expected: Action, envelope: &Envelope) -> Result<(), KpxError> {
    if envelope.action == expected.as_str() {
        Ok(())
    } else {
        Err(KpxError::UnexpectedAction {
            expected: expected.as_str().to_string(),
            received: envelope.action.clone(),
        })
    }
}

fn parse_success(payload: &Value) -> Option<bool> {
    let success_value = payload.get("success")?;
    let success: SuccessValue = serde_json::from_value(success_value.clone()).ok()?;
    Some(success.is_true())
}

#[allow(deprecated)]
fn make_nonce(bytes: &[u8; NONCE_LEN]) -> GenericArray<u8, <SalsaBox as AeadCore>::NonceSize> {
    GenericArray::clone_from_slice(bytes)
}

fn supports_identification_key(version: Option<&str>) -> bool {
    const MIN_SUPPORTED: (u32, u32, u32) = (2, 3, 4);
    version
        .and_then(parse_version_triplet)
        .map(|parsed| parsed >= MIN_SUPPORTED)
        .unwrap_or(true)
}

fn parse_version_triplet(input: &str) -> Option<(u32, u32, u32)> {
    let core = input.split(['-', '+']).next()?.trim();
    if core.is_empty() {
        return None;
    }

    let mut parts = core.split('.');
    let major = parts.next()?.parse().ok()?;
    let minor = parts.next().unwrap_or("0").parse().ok()?;
    let patch = parts.next().unwrap_or("0").parse().ok()?;

    Some((major, minor, patch))
}

#[instrument(level = "trace")]
fn generate_keypair() -> (PublicKey, SecretKey) {
    let mut rng = OsRng;
    let secret = SecretKey::generate(&mut rng);
    let public = secret.public_key();
    (public, secret)
}

struct RequestBuilder {
    map: Map<String, Value>,
}

impl RequestBuilder {
    fn new(action: Action) -> Self {
        let mut map = Map::new();
        map.insert("action".into(), Value::String(action.as_str().to_string()));
        Self { map }
    }

    fn with_client_id(mut self, client_id: &str) -> Self {
        self.map
            .insert("clientID".into(), Value::String(client_id.to_string()));
        self
    }

    fn with_nonce(mut self, nonce: &[u8; NONCE_LEN]) -> Self {
        self.map
            .insert("nonce".into(), Value::String(encode_base64(nonce)));
        self
    }

    fn with_value(mut self, key: &str, value: Value) -> Self {
        self.map.insert(key.into(), value);
        self
    }

    fn extend(mut self, extra: Map<String, Value>) -> Self {
        self.map.extend(extra);
        self
    }

    fn build(self) -> Map<String, Value> {
        self.map
    }
}

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

    fn decode_sent_json(sent: &str) -> Value {
        serde_json::from_str(sent).expect("sent json")
    }

    #[test]
    fn change_public_keys_stores_host_key() {
        let (host_public, _) = generate_keypair();
        let host_public_b64 = encode_base64(host_public.as_bytes());
        let response = json!({
            "action": "change-public-keys",
            "publicKey": host_public_b64,
            "success": "true",
            "version": "2.7.0"
        })
        .to_string();

        let transport = MockTransport::with_responses(vec![response]);
        let mut client = KeePassXcClient::new(transport.clone());
        let _ = client.change_public_keys().unwrap();
        assert!(client.host_public_key.is_some());

        let sent = transport.sent.lock().unwrap();
        let request = decode_sent_json(&sent[0]);
        assert_eq!(
            request.get("action").and_then(Value::as_str).unwrap(),
            "change-public-keys"
        );
        assert!(request.get("publicKey").is_some());
        assert!(request.get("clientID").is_some());
    }

    #[test]
    fn associate_adds_record() {
        let (host_public, host_secret) = generate_keypair();
        let host_public_b64 = encode_base64(host_public.as_bytes());

        let handshake_response = json!({
            "action": "change-public-keys",
            "publicKey": host_public_b64,
            "success": "true",
            "version": "2.7.0"
        })
        .to_string();

        let transport = MockTransport::with_responses(vec![handshake_response.clone()]);
        let mut client = KeePassXcClient::new(transport.clone());
        client.change_public_keys().unwrap();

        // Inspect the handshake request to recover client's public key
        let client_public_b64 = {
            let sent = transport.sent.lock().unwrap();
            decode_sent_json(&sent[0])
                .get("publicKey")
                .and_then(Value::as_str)
                .unwrap()
                .to_string()
        };
        let client_public_bytes = decode_base64_array::<KEY_LEN>(&client_public_b64).unwrap();
        let client_public = PublicKey::from(client_public_bytes);

        // Prepare response payload
        let response_payload = json!({
            "hash": "29234e32274a32276e25666a42",
            "version": "2.7.0",
            "success": "true",
            "id": "testclient",
            "nonce": "nonce"
        });

        let nonce_bytes = random_bytes::<NONCE_LEN>();
        let nonce = make_nonce(&nonce_bytes);
        let nonce_b64 = encode_base64(&nonce_bytes);
        let salsa = SalsaBox::new(&client_public, &host_secret);
        let cipher = salsa
            .encrypt(&nonce, response_payload.to_string().as_bytes())
            .unwrap();
        let response_json = json!({
            "action": "associate",
            "nonce": nonce_b64,
            "message": encode_base64(&cipher)
        })
        .to_string();
        transport.push_response(response_json);

        let record = client.associate().unwrap();
        assert_eq!(record.association.id, "testclient");
        assert_eq!(record.database_hash, "29234e32274a32276e25666a42");

        assert_eq!(client.associations.len(), 1);
    }

    #[test]
    fn associate_uses_session_key_for_legacy_versions() {
        let (host_public, host_secret) = generate_keypair();
        let host_public_b64 = encode_base64(host_public.as_bytes());

        let handshake_response = json!({
            "action": "change-public-keys",
            "publicKey": host_public_b64,
            "success": "true",
            "version": "2.3.3"
        })
        .to_string();

        let transport = MockTransport::with_responses(vec![handshake_response.clone()]);
        let mut client = KeePassXcClient::new(transport.clone());
        client.change_public_keys().unwrap();

        // Capture the client's session key from the handshake payload.
        let client_public_b64 = {
            let sent = transport.sent.lock().unwrap();
            decode_sent_json(&sent[0])
                .get("publicKey")
                .and_then(Value::as_str)
                .unwrap()
                .to_string()
        };
        let client_public_bytes = decode_base64_array::<KEY_LEN>(&client_public_b64).unwrap();
        let client_public = PublicKey::from(client_public_bytes);

        let response_payload = json!({
            "hash": "legacy-hash",
            "version": "2.3.3",
            "success": "true",
            "id": "legacy-client",
            "nonce": "nonce"
        });

        let nonce_bytes = random_bytes::<NONCE_LEN>();
        let salsa = SalsaBox::new(&client_public, &host_secret);
        let cipher = salsa
            .encrypt(
                &make_nonce(&nonce_bytes),
                response_payload.to_string().as_bytes(),
            )
            .unwrap();
        transport.push_response(
            json!({
                "action": "associate",
                "nonce": encode_base64(&nonce_bytes),
                "message": encode_base64(&cipher)
            })
            .to_string(),
        );

        let record = client.associate().unwrap();
        assert_eq!(record.association.key, client_public_b64);
    }

    #[test]
    fn get_logins_decrypts_entries() {
        let (host_public, host_secret) = generate_keypair();
        let host_public_b64 = encode_base64(host_public.as_bytes());
        let handshake_response = json!({
            "action": "change-public-keys",
            "publicKey": host_public_b64,
            "success": "true",
            "version": "2.7.0"
        })
        .to_string();

        let transport = MockTransport::with_responses(vec![handshake_response]);
        let mut client = KeePassXcClient::new(transport.clone());
        client.change_public_keys().unwrap();

        let client_public_b64 = {
            let sent = transport.sent.lock().unwrap();
            decode_sent_json(&sent[0])
                .get("publicKey")
                .and_then(Value::as_str)
                .unwrap()
                .to_string()
        };
        let client_public_bytes = decode_base64_array::<KEY_LEN>(&client_public_b64).unwrap();
        let client_public = PublicKey::from(client_public_bytes);

        let assoc_payload = json!({
            "hash": "hash",
            "version": "2.7.0",
            "success": "true",
            "id": "testclient",
            "nonce": "nonce"
        });
        let nonce_assoc = random_bytes::<NONCE_LEN>();
        let salsa = SalsaBox::new(&client_public, &host_secret);
        let cipher_assoc = salsa
            .encrypt(
                &make_nonce(&nonce_assoc),
                assoc_payload.to_string().as_bytes(),
            )
            .unwrap();
        transport.push_response(
            json!({
                "action": "associate",
                "nonce": encode_base64(&nonce_assoc),
                "message": encode_base64(&cipher_assoc),
            })
            .to_string(),
        );
        let association = client.associate().unwrap();

        let logins_payload = json!({
            "count": "1",
            "entries": [{
                "login": "user",
                "name": "user",
                "password": "secret"
            }],
            "success": "true",
            "nonce": "nonce",
            "version": "2.7.0"
        });
        let nonce_logins = random_bytes::<NONCE_LEN>();
        let cipher_logins = salsa
            .encrypt(
                &make_nonce(&nonce_logins),
                logins_payload.to_string().as_bytes(),
            )
            .unwrap();
        transport.push_response(
            json!({
                "action": "get-logins",
                "nonce": encode_base64(&nonce_logins),
                "message": encode_base64(&cipher_logins),
            })
            .to_string(),
        );

        let query = LoginQuery::new("https://example.com", vec![association.association.clone()])
            .with_primary_id(association.association.id.clone());
        let response = client.get_logins(&query).unwrap();
        let entries = response.entries.unwrap();
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].login, "user");
        assert_eq!(entries[0].password, "secret");

        let sent = transport.sent.lock().unwrap();
        let request = decode_sent_json(&sent[2]);
        let nonce_b64 = request.get("nonce").and_then(Value::as_str).unwrap();
        let message_b64 = request.get("message").and_then(Value::as_str).unwrap();
        let nonce_bytes = decode_base64_array::<NONCE_LEN>(nonce_b64).unwrap();
        let cipher = decode_base64(message_b64).unwrap();
        let salsa = SalsaBox::new(&client_public, &host_secret);
        let plaintext = salsa
            .decrypt(&make_nonce(&nonce_bytes), cipher.as_slice())
            .unwrap();
        let payload: Value = serde_json::from_slice(&plaintext).unwrap();
        assert_eq!(
            payload.get("id").and_then(Value::as_str),
            Some("testclient")
        );
    }

    #[test]
    fn get_totp_returns_code() {
        let (host_public, host_secret) = generate_keypair();
        let host_public_b64 = encode_base64(host_public.as_bytes());
        let handshake_response = json!({
            "action": "change-public-keys",
            "publicKey": host_public_b64,
            "success": "true",
            "version": "2.7.0"
        })
        .to_string();

        let transport = MockTransport::with_responses(vec![handshake_response]);
        let mut client = KeePassXcClient::new(transport.clone());
        client.change_public_keys().unwrap();

        let totp_payload = json!({
            "totp": "123456",
            "success": "true",
            "version": "2.7.7"
        });
        let nonce_bytes = random_bytes::<NONCE_LEN>();
        let salsa = SalsaBox::new(&client.public_key, &host_secret);
        let cipher = salsa
            .encrypt(
                &make_nonce(&nonce_bytes),
                totp_payload.to_string().as_bytes(),
            )
            .unwrap();
        transport.push_response(
            json!({
                "action": "get-totp",
                "nonce": encode_base64(&nonce_bytes),
                "message": encode_base64(&cipher)
            })
            .to_string(),
        );

        let totp = client.get_totp("entry-uuid").unwrap();
        assert_eq!(totp, Some("123456".to_string()));

        let sent = transport.sent.lock().unwrap();
        let request = decode_sent_json(&sent[1]);
        assert_eq!(
            request.get("action").and_then(Value::as_str),
            Some("get-totp")
        );
    }
}