agentknock 0.1.0

use std::{collections::BTreeMap, future::Future};

use serde::{Deserialize, Serialize};
use thiserror::Error;
use ulid::Ulid;

use crate::{
    Client, RequestError,
    config::{ConfigurationError, clear_rotation_key, read_pairing_from},
    crypto::Session,
    pairing::RotationError,
    protocol::{self, Method, Response},
    websocket::RelayExchange,
};

/// Metadata for a secret available from the paired device.
///
/// This type never contains secret values.
#[derive(Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum Secret {
    /// A secret that provides environment variables.
    #[non_exhaustive]
    Environment {
        /// An optional human-readable description.
        description: Option<String>,
        /// The environment variable names provided by the secret.
        ///
        /// Names are sorted and contain no duplicates.
        variables: Vec<String>,
    },
}

/// Secret metadata keyed by secret name.
///
/// Iteration yields secret names in lexicographic order.
pub type Secrets = BTreeMap<String, Secret>;

/// An environment-variable secret to upload to the device.
#[derive(Debug, Eq, PartialEq)]
pub struct EnvironmentSecret {
    /// The name of the new or existing secret.
    ///
    /// In [`SecretUploadMode::Create`] the device can let the user choose a
    /// different name when accepting the upload. In the other modes, this must
    /// identify the existing secret to change.
    pub name: String,

    /// The proposed description.
    ///
    /// In [`SecretUploadMode::Update`], `None` retains the existing
    /// description. An empty string proposes removing it.
    pub description: Option<String>,

    /// Environment variable values keyed by variable name.
    pub variables: BTreeMap<String, String>,
}

/// How an uploaded secret changes device state after user acceptance.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SecretUploadMode {
    /// Proposes a new secret whose final name the user can choose on the device.
    Create,

    /// Replaces an existing secret, removing variables that aren't supplied.
    Replace,

    /// Updates supplied fields while retaining variables that aren't supplied.
    Update,
}

/// A stage reported while a secret-list request is running.
///
/// A successful operation reports `Preparing`, `WaitingForDelivery`,
/// optionally one or more `WaitingForResponse` updates, `Completing`, and
/// `Completed`, in that order. An operation that fails stops without reporting
/// `Completed`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum SecretListProgress {
    /// Agentknock is reading local state and preparing the protected request.
    Preparing,

    /// The request is waiting to be delivered to the device.
    WaitingForDelivery,

    /// The device has received the request but hasn't returned a response.
    WaitingForResponse,

    /// Agentknock is validating the response and handing off the completion.
    Completing,

    /// The operation has finished successfully.
    Completed,
}

/// A stage reported while a secret upload is running.
///
/// A completed exchange reports `Preparing`, `WaitingForDelivery`, optionally
/// one or more `WaitingForResponse` updates, `Completing`, and `Completed`, in
/// that order. A rejected upload can return an error after reporting
/// `Completed`; other failures stop without reporting `Completed`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum SecretUploadProgress {
    /// Agentknock is reading local state and preparing the protected upload.
    Preparing,

    /// The upload is waiting to be delivered to the device.
    WaitingForDelivery,

    /// The device has received the upload but hasn't confirmed receipt.
    WaitingForResponse,

    /// Agentknock is validating the response and handing off the completion.
    Completing,

    /// The device confirmed receipt and the operation has finished.
    Completed,
}

/// An error uploading a secret.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum SecretUploadError {
    /// The authenticated exchange failed.
    #[error(transparent)]
    Request(#[from] RequestError),

    /// The device rejected the upload instead of storing it for user review.
    #[error("the device rejected the secret upload: {message}")]
    Rejected {
        /// Human-readable context supplied by the device.
        message: String,
    },
}

impl From<ConfigurationError> for SecretUploadError {
    fn from(error: ConfigurationError) -> Self {
        Self::Request(error.into())
    }
}

impl From<crate::websocket::Error> for SecretUploadError {
    fn from(error: crate::websocket::Error) -> Self {
        Self::Request(error.into())
    }
}

impl Client {
    /// Lists metadata for secrets available from the paired device.
    ///
    /// The returned [`Secrets`] includes names, types, descriptions, and the
    /// names of values each secret provides. It never includes secret values.
    ///
    /// The `progress` callback receives lifecycle updates synchronously and
    /// should return promptly. Cancellation before an authenticated response
    /// returns [`RequestError::Interrupted`]. After a response is
    /// authenticated, cancellation only shortens the best-effort completion
    /// handoff and the method still returns the decoded list. Pass
    /// [`std::future::pending()`] when the operation doesn't need cancellation.
    ///
    /// # Errors
    ///
    /// Returns [`RequestError`] if local pairing state isn't active, the relay
    /// exchange fails, the response is invalid, or the operation is canceled
    /// before a response is authenticated.
    pub async fn list_secrets<P>(
        &self,
        cancellation: impl Future<Output = ()>,
        mut progress: P,
    ) -> Result<Secrets, RequestError>
    where
        P: FnMut(SecretListProgress),
    {
        tokio::pin!(cancellation);
        progress(SecretListProgress::Preparing);
        self.prepare_request()?;
        let pairing_path = self.pairing_path()?;
        let pairing = read_pairing_from(&pairing_path)?;
        let request_id = Ulid::generate();
        let plaintext = self
            .encode(&ListRequest {
                method: Method::SecretList,
            })
            .map_err(RequestError::other)?;
        let mut session = Session::new(&pairing, &request_id).map_err(RequestError::other)?;
        let request = session
            .seal_request(&plaintext)
            .map_err(RequestError::other)?;
        let mut relay = RelayExchange::authenticated(&pairing, &request_id.to_string())?;

        progress(SecretListProgress::WaitingForDelivery);
        let response = tokio::select! {
            biased;
            _ = cancellation.as_mut() => return Err(RequestError::Interrupted),
            response = relay.request(&request, || {
                progress(SecretListProgress::WaitingForResponse);
            }) => response?,
        };
        progress(SecretListProgress::Completing);
        let plaintext = session
            .open_response(response)
            .map_err(RequestError::other)?;
        if let Some(rotation_key) = pairing.rotation_key() {
            clear_rotation_key(&pairing_path, rotation_key)?;
        }
        let response: ListResponse =
            match protocol::decode_response(&plaintext).map_err(RequestError::other)? {
                Response::Message(response) => response,
                Response::Error(error) => {
                    if let Some(completion) =
                        protocol::seal_error_completion(self, &mut session, &error)
                    {
                        let _ = relay.complete_briefly(&completion).await;
                    }
                    return Err(RequestError::DeviceRejected {
                        code: error.code,
                        message: error.message,
                    });
                }
            };
        let plaintext = self.encode(&EmptyMessage {}).map_err(RequestError::other)?;
        let completion = session
            .seal_completion(&plaintext)
            .map_err(RequestError::other)?;
        let interrupted = tokio::select! {
            biased;
            _ = cancellation.as_mut() => true,
            result = relay.complete(&completion) => {
                result?;
                false
            }
        };
        if interrupted {
            let _ = relay.complete_briefly(&completion).await;
        }
        progress(SecretListProgress::Completed);

        Ok(response
            .secrets
            .into_iter()
            .map(|(name, secret)| (name, secret.into()))
            .collect())
    }

    /// Uploads an environment-variable secret for review on the device.
    ///
    /// Success means that the device received and stored the upload proposal.
    /// It doesn't mean that the user accepted the proposal or that the secret
    /// is available for use. Upload mode controls how a later acceptance would
    /// change device state.
    ///
    /// The `progress` callback receives lifecycle updates synchronously and
    /// should return promptly. Cancellation before an authenticated response
    /// returns [`RequestError::Interrupted`] through [`SecretUploadError`].
    /// After a response is authenticated, cancellation only shortens the
    /// best-effort completion handoff and the method still returns the device's
    /// result. Pass [`std::future::pending()`] when the operation doesn't need
    /// cancellation.
    ///
    /// # Errors
    ///
    /// Returns [`SecretUploadError`] if local pairing state isn't active, the
    /// relay exchange fails, the device rejects the proposal, the response is
    /// invalid, or the operation is canceled before a response is
    /// authenticated.
    pub async fn upload_secret<P>(
        &self,
        secret: &EnvironmentSecret,
        mode: SecretUploadMode,
        cancellation: impl Future<Output = ()>,
        mut progress: P,
    ) -> Result<(), SecretUploadError>
    where
        P: FnMut(SecretUploadProgress),
    {
        tokio::pin!(cancellation);
        progress(SecretUploadProgress::Preparing);
        self.prepare_request()?;
        let pairing_path = self.pairing_path()?;
        let pairing = read_pairing_from(&pairing_path)?;
        let request_id = Ulid::generate();
        let request_payload = UploadRequest {
            method: Method::SecretUpload,
            mode: mode.into(),
            secret: NamedSecretMessage::from(secret),
        };
        let plaintext = self.encode(&request_payload).map_err(RequestError::other)?;
        let mut session = Session::new(&pairing, &request_id).map_err(RequestError::other)?;
        let request = session
            .seal_request(&plaintext)
            .map_err(RequestError::other)?;
        let mut relay = RelayExchange::authenticated(&pairing, &request_id.to_string())?;

        progress(SecretUploadProgress::WaitingForDelivery);
        let response = tokio::select! {
            biased;
            _ = cancellation.as_mut() => {
                return Err(RequestError::Interrupted.into());
            }
            response = relay.request(&request, || {
                progress(SecretUploadProgress::WaitingForResponse);
            }) => response?,
        };
        progress(SecretUploadProgress::Completing);
        let plaintext = session
            .open_response(response)
            .map_err(RequestError::other)?;
        if let Some(rotation_key) = pairing.rotation_key() {
            clear_rotation_key(&pairing_path, rotation_key)?;
        }
        let response: UploadResult =
            match protocol::decode_response(&plaintext).map_err(RequestError::other)? {
                Response::Message(response) => response,
                Response::Error(error) => {
                    if let Some(completion) =
                        protocol::seal_error_completion(self, &mut session, &error)
                    {
                        let _ = relay.complete_briefly(&completion).await;
                    }
                    return Err(RequestError::DeviceRejected {
                        code: error.code,
                        message: error.message,
                    }
                    .into());
                }
            };
        let completion = self.encode(&response).map_err(RequestError::other)?;
        let completion = session
            .seal_completion(&completion)
            .map_err(RequestError::other)?;
        tokio::select! {
            biased;
            _ = cancellation.as_mut() => {},
            _ = relay.complete_briefly(&completion) => {},
        }
        progress(SecretUploadProgress::Completed);

        match response {
            UploadResult::Received => Ok(()),
            UploadResult::Rejected { message } => Err(SecretUploadError::Rejected { message }),
        }
    }

    fn prepare_request(&self) -> Result<(), RequestError> {
        self.maybe_rotate_psk().map_err(|error| match error {
            RotationError::Configuration(error) => RequestError::Configuration(error),
            RotationError::Other(error) => RequestError::Other(error),
        })?;
        Ok(())
    }
}

#[derive(Serialize)]
struct ListRequest {
    method: Method,
}

#[derive(Serialize)]
struct EmptyMessage {}

#[derive(Deserialize)]
struct ListResponse {
    secrets: BTreeMap<String, SecretMessage<Vec<String>>>,
}

#[derive(Serialize)]
struct UploadRequest {
    method: Method,
    mode: SecretUploadModeMessage,
    secret: NamedSecretMessage<BTreeMap<String, EnvironmentVariableMessage>>,
}

#[derive(Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
enum SecretUploadModeMessage {
    Create,
    Replace,
    Update,
}

impl From<SecretUploadMode> for SecretUploadModeMessage {
    fn from(mode: SecretUploadMode) -> Self {
        match mode {
            SecretUploadMode::Create => Self::Create,
            SecretUploadMode::Replace => Self::Replace,
            SecretUploadMode::Update => Self::Update,
        }
    }
}

#[derive(Deserialize, Serialize)]
#[serde(tag = "result", rename_all = "SCREAMING_SNAKE_CASE")]
enum UploadResult {
    Received,
    Rejected { message: String },
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SecretMessage<T> {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) description: Option<String>,
    #[serde(flatten)]
    pub(crate) contents: SecretContentsMessage<T>,
}

#[derive(Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(crate) enum SecretContentsMessage<T> {
    Environment { variables: T },
}

#[derive(Deserialize, Serialize)]
pub(crate) struct EnvironmentVariableMessage {
    pub(crate) value: String,
}

#[derive(Serialize)]
struct NamedSecretMessage<T> {
    name: String,
    #[serde(flatten)]
    secret: SecretMessage<T>,
}

impl From<&EnvironmentSecret> for NamedSecretMessage<BTreeMap<String, EnvironmentVariableMessage>> {
    fn from(secret: &EnvironmentSecret) -> Self {
        Self {
            name: secret.name.clone(),
            secret: SecretMessage {
                description: secret.description.clone(),
                contents: SecretContentsMessage::Environment {
                    variables: secret
                        .variables
                        .iter()
                        .map(|(name, value)| {
                            (
                                name.clone(),
                                EnvironmentVariableMessage {
                                    value: value.clone(),
                                },
                            )
                        })
                        .collect(),
                },
            },
        }
    }
}

impl From<SecretMessage<Vec<String>>> for Secret {
    fn from(secret: SecretMessage<Vec<String>>) -> Self {
        match secret.contents {
            SecretContentsMessage::Environment { mut variables } => {
                variables.sort();
                variables.dedup();
                Self::Environment {
                    description: secret.description,
                    variables,
                }
            }
        }
    }
}