matroskin 0.1.0

//! Define Auth Data module
//!
//! - Item: [AuthData]
//! - ApiDoc: <https://apidoc.whatsminer.com/#api-Token-generate_token>
use aes::Aes256;
use aes::cipher::{BlockEncrypt, KeyInit, block_padding::Pkcs7};
use base64::{Engine, engine::general_purpose::STANDARD};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::fmt::Debug;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::trace;
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

use crate::account::Account;
use crate::command::Command;
use crate::error::{Error, Result};
use crate::redacted::Redacted;

#[derive(Debug, Clone, PartialEq, Serialize, Zeroize, ZeroizeOnDrop)]
/// Authentication data required for secured/encrypted command endpoints
pub struct AuthData {
    /// UNIX timestamp when the command was issued
    pub ts: u64,
    /// Auth miner token
    ///
    /// It generated from 'command'+'account's password'+'salt'+'ts'
    ///
    /// - ApiDoc: <https://apidoc.whatsminer.com/#api-Token-generate_token>
    #[serde(serialize_with = "serialize_token")]
    token: Redacted<[u8; 8]>,
    /// Account username
    #[zeroize(skip)]
    #[serde(rename = "account")]
    username: Account,
    /// SHA256 raw bytes for encrypting params (AES-256 key)
    #[serde(skip)]
    aes_key: Redacted<[u8; 32]>,
}

fn serialize_token<S>(
    token: &Redacted<[u8; 8]>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    let s = std::str::from_utf8(&**token).map_err(serde::ser::Error::custom)?;
    serializer.serialize_str(s)
}

impl AuthData {
    /// Generate authentication for the command using the current Unix timestamp.
    pub fn new<C: Command>(
        username: Account,
        password: impl AsRef<str>,
        salt: &str,
    ) -> Result<Self> {
        Self::for_command_name(C::CMD_NAME, username, password, salt)
    }

    /// Generate auth data for an explicit command name
    pub fn for_command_name(
        cmd_name: &str,
        username: Account,
        password: impl AsRef<str>,
        salt: &str,
    ) -> Result<Self> {
        let ts = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
        let mut timestamp = itoa::Buffer::new();
        let mut hash = Sha256::new();
        hash.update(cmd_name.as_bytes());
        hash.update(password.as_ref().as_bytes());
        hash.update(salt.as_bytes());
        hash.update(timestamp.format(ts).as_bytes());
        let aes_key = Redacted(<[u8; 32]>::from(hash.finalize()));
        let mut token = Redacted([0; 8]);
        // Six digest bytes encode to exactly the eight token characters, without padding.
        STANDARD
            .encode_slice(&aes_key[..6], &mut token.0)
            .map_err(|_| Error::CorruptedToken)?;

        trace!(command = %cmd_name, ts, "Auth data generated.");

        Ok(Self {
            ts,
            token,
            username,
            aes_key,
        })
    }

    /// Encrypt data using AES-256-ECB with PKCS7 padding
    ///
    /// According to WhatsMiner API documentation, params are encrypted using:
    /// - AES-256 in ECB mode (no IV needed)
    /// - PKCS7 padding
    /// - The SHA256 hash as the encryption key
    pub fn encrypt(&self, data: impl AsRef<[u8]>) -> Result<String> {
        let plaintext = data.as_ref();

        let cipher = Aes256::new_from_slice(&*self.aes_key).map_err(|_| Error::EncryptionFailed)?;
        let padded_len = plaintext
            .len()
            .checked_add(16 - plaintext.len() % 16)
            .ok_or(Error::EncryptionFailed)?;
        let mut buf = Zeroizing::new(Vec::with_capacity(padded_len));
        buf.extend_from_slice(plaintext);
        buf.resize(padded_len, 0);
        let encrypted = cipher
            .encrypt_padded::<Pkcs7>(&mut buf, plaintext.len())
            .map_err(|_| Error::EncryptionFailed)?;

        let encrypted_len = encrypted.len();
        trace!(
            plaintext_len = data.as_ref().len(),
            ciphertext_len = encrypted_len,
            "Params encrypted."
        );

        Ok(STANDARD.encode(encrypted))
    }
}

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

    #[test]
    fn token_matches_protocol_digest_prefix() {
        let auth =
            AuthData::for_command_name("set.miner.pools", Account::Super, "secret", "QbVy1Ou3")
                .unwrap();
        let json = serde_json::to_value(&auth).unwrap();
        let timestamp = json["ts"].as_u64().unwrap();
        let expected =
            Sha256::digest(format!("set.miner.poolssecretQbVy1Ou3{timestamp}").as_bytes());
        let base64 = base64_light::base64_encode_bytes(&expected);
        assert_eq!(json["token"], &base64[..8]);
        assert_eq!(json["account"], "super");
        assert!(json.get("aes_key").is_none());
        let mut ciphertext =
            base64_light::base64_decode(&auth.encrypt(b"{\"enabled\":true}").unwrap());
        let cipher = Aes256::new_from_slice(&expected).unwrap();
        let cleartext = cipher.decrypt_padded_mut::<Pkcs7>(&mut ciphertext).unwrap();
        assert_eq!(cleartext, b"{\"enabled\":true}");
    }

    #[test]
    fn debug_redacts_authentication_material() {
        let auth = AuthData::for_command_name("cmd", Account::Super, "password", "salt").unwrap();
        let debug = format!("{auth:?}");
        assert!(!debug.contains(std::str::from_utf8(&*auth.token).unwrap()));
        assert!(!debug.contains("password"));
        assert!(!debug.contains("salt"));
    }
}