//! SSH Agent client
//!
//! The `sshca-cli user login` command will automatically add the
//! signed certificate it received from SSHCA to the user's SSH agent.

use ssh_encoding::{Encode, Writer};
use ssh_key::{private::KeypairData, Certificate, PrivateKey};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixStream;
use tracing::{debug, info, trace};

/// Error type for SSH agent communication
#[derive(Debug, thiserror::Error)]
pub enum AgentError {
    #[error("SSH agent not found: {0}")]
    NotFound(std::env::VarError),
    #[error("Could not connect to SSH agent: {0}")]
    Connect(std::io::Error),
    #[error("Encoding error: {0}")]
    Encoding(#[from] ssh_encoding::Error),
    #[error("Invalid message length: {0}")]
    IvalidMessageLength(#[from] std::num::TryFromIntError),
    #[error("Error communicating with SSH agent: {0}")]
    Io(#[from] std::io::Error),
    #[error("SSH agent returned failure")]
    Failure,
}

/// SSH Agent message types
#[repr(u8)]
enum AgentMessageType {
    AddIdentity = 17,
}

/// SSH Agent response code
enum AgentResponseCode {
    Success,
    Failure,
    InvalidFormat,
}

impl From<u8> for AgentResponseCode {
    fn from(v: u8) -> Self {
        // https://github.com/openssh/openssh-portable/blob/V_9_0_P1/authfd.c#L74
        match v {
            6 => Self::Success,
            30 | 102 | 229 => Self::Failure,
            _ => Self::InvalidFormat,
        }
    }
}

/// Add an SSH certificate and private key to the SSH Agent
///
/// This function adds an SSH certificate and its corresponding private
/// key to the SSH Agent.
///
/// If an error occurs while attempting to add the key to the agent,
/// including if no agent is configured, [`AgentError`] is returned.
pub async fn add_to_agent(
    key: &PrivateKey,
    cert: &Certificate,
) -> Result<(), AgentError> {
    let sock_path =
        std::env::var("SSH_AUTH_SOCK").map_err(AgentError::NotFound)?;
    debug!("Connecting to SSH agent at {:?}", sock_path);
    let mut sock = UnixStream::connect(sock_path)
        .await
        .map_err(AgentError::Connect)?;

    trace!("Serializing SSH2_AGENTC_ADD_IDENTITY message");
    let mut buf: Vec<u8> = vec![];
    buf.push(AgentMessageType::AddIdentity as u8);
    serialize_key_cert(key, cert, &mut buf)?;
    let len = u32::try_from(buf.len())?;

    debug!("Sending key to SSH agent");
    sock.write_all(&len.to_be_bytes()).await?;
    sock.write_all(&buf[..]).await?;

    let mut res_len = [0u8; 4];
    trace!("Waiting for SSH agent response");
    sock.read_exact(&mut res_len[..]).await?;
    let res_len = usize::try_from(u32::from_be_bytes(res_len))?;
    let mut res = vec![0u8; res_len];
    trace!("Reading {} bytes from SSH agent", res_len);
    sock.read_exact(&mut res).await?;
    trace!("Received SSH agent response: {:?}", res);

    match res[0].into() {
        AgentResponseCode::Success => {
            info!("Successfully added SSH user certificate to SSH Agent");
            Ok(())
        }
        _ => Err(AgentError::Failure),
    }
}

/// Serialize an SSH certificate and private key to send to SSH Agent
///
/// This function takes a byte buffer and fills it with the wire
/// representation of an SSH certificate and its corresponding provate
/// key, in order to send them to the SSH Agent.
///
/// The [draft-miller-ssh-agent-11][0] protocol does not specify how
/// certificates are sent to the SSH Agent.  The message format used
/// here was discovered by reading the OpenSSH portable code,
/// specifically [sshkey.c][1], and observing communications between
/// `ssh-add` and `ssh-agent`.
///
/// [0]: https://datatracker.ietf.org/doc/html/draft-miller-ssh-agent-11
/// [1]: https://github.com/openssh/openssh-portable/blob/V_9_0_P1/sshkey.c#L3230
fn serialize_key_cert(
    key: &PrivateKey,
    cert: &Certificate,
    buf: &mut impl Writer,
) -> Result<(), ssh_encoding::Error> {
    cert.algorithm().to_certificate_type().encode(buf)?;
    cert.encode_prefixed(buf)?;
    match key.key_data() {
        KeypairData::Dsa(k) => k.encode(buf)?,
        KeypairData::Ecdsa(k) => k.encode(buf)?,
        KeypairData::Ed25519(k) => k.encode(buf)?,
        KeypairData::Encrypted(k) => k.encode(buf)?,
        KeypairData::Rsa(k) => k.encode(buf)?,
        KeypairData::SkEcdsaSha2NistP256(k) => k.encode(buf)?,
        KeypairData::SkEd25519(k) => k.encode(buf)?,
        KeypairData::Other(k) => k.encode(buf)?,
        &_ => todo!(),
    };
    key.comment().encode(buf)?;
    Ok(())
}