rust-async-tuyapi 1.2.0

use aes::cipher::block_padding::Pkcs7;
use base64::Engine;
use hmac::{Hmac, Mac};
use sha2::Sha256;

use crate::mesparse::TuyaVersion;
use crate::Result;

/// TuyaCipher is a low level api for encrypting and decrypting Vec<u8>'s.
#[derive(Clone)]
pub(crate) struct TuyaCipher {
    key: Vec<u8>,
    version: TuyaVersion,
}

type Aes128EcbEnc = ecb::Encryptor<aes::Aes128>;
type Aes128EcbDec = ecb::Decryptor<aes::Aes128>;
type HmacSha256 = Hmac<Sha256>;

fn maybe_strip_header(version: &TuyaVersion, data: &[u8]) -> Vec<u8> {
    if data.len() > 3 && &data[..3] == version.as_bytes() {
        match version {
            TuyaVersion::ThreeOne => data.split_at(19).1.to_vec(),
            TuyaVersion::ThreeThree | TuyaVersion::ThreeFour => data.split_at(15).1.to_vec(),
        }
    } else {
        data.to_vec()
    }
}

impl TuyaCipher {
    pub fn create(key: &[u8], version: TuyaVersion) -> TuyaCipher {
        TuyaCipher {
            key: key.to_vec(),
            version,
        }
    }

    pub fn set_key(&mut self, key: Vec<u8>) {
        self.key = key
    }

    pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
        use aes::cipher::{BlockEncryptMut, KeyInit};

        let ct = Aes128EcbEnc::new_from_slice(self.key.as_slice())?
            .encrypt_padded_vec_mut::<Pkcs7>(data);

        match self.version {
            TuyaVersion::ThreeOne => Ok(base64::engine::general_purpose::STANDARD
                .encode(ct)
                .as_bytes()
                .to_vec()),
            TuyaVersion::ThreeThree | TuyaVersion::ThreeFour => Ok(ct.to_vec()),
        }
    }

    pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>> {
        use aes::cipher::{BlockDecryptMut, KeyInit};

        // Different header size in version 3.1 and 3.3
        let data = maybe_strip_header(&self.version, data);

        // 3.1 is base64 encoded, 3.3 is not
        let data = match self.version {
            TuyaVersion::ThreeOne => base64::engine::general_purpose::STANDARD.decode(&data)?,
            TuyaVersion::ThreeThree | TuyaVersion::ThreeFour => data.to_vec(),
        };

        let pt = Aes128EcbDec::new_from_slice(self.key.as_slice())?
            .decrypt_padded_vec_mut::<Pkcs7>(&data)?;

        Ok(pt.to_vec())
    }

    pub fn md5(&self, payload: &[u8]) -> Vec<u8> {
        let hash_line: Vec<u8> = [
            b"data=",
            payload,
            b"||lpv=",
            self.version.as_bytes(),
            b"||",
            self.key.as_ref(),
        ]
        .iter()
        .flat_map(|bytes| bytes.iter())
        .copied()
        .collect();
        let digest: [u8; 16] = md5::compute(hash_line).into();
        digest[4..16].to_vec()
    }

    pub fn hmac(&self, payload: &[u8]) -> Result<Vec<u8>> {
        let mut mac = HmacSha256::new_from_slice(&self.key)?;
        mac.update(payload);
        let result = mac.finalize();
        Ok(result.into_bytes().to_vec())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn maybe_strip_header_with_correct_header() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let message = b"3.133ed3d4a21effe90zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=";
        let expected = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=".to_vec();
        assert_eq!(maybe_strip_header(&cipher.version, message), expected)
    }

    #[test]
    fn maybe_strip_header_without_header() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let message = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=".to_vec();
        assert_eq!(maybe_strip_header(&cipher.version, &message), message)
    }

    #[test]
    fn encrypt_message() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let data =
            r#"{"devId":"002004265ccf7fb1b659","dps":{"1":false,"2":0},"t":1529442366,"s":8}"#
                .as_bytes();
        let result = cipher.encrypt(data).unwrap();

        let expected = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=".to_vec();
        assert_eq!(expected, result);
    }

    #[test]
    fn encrypt_message_without_base64_encoding() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let data =
            r#"{"devId":"002004265ccf7fb1b659","dps":{"1":false,"2":0},"t":1529442366,"s":8}"#
                .as_bytes();
        let result = cipher.encrypt(data).unwrap();

        let expected = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=".to_vec();
        assert_eq!(expected, result);
    }

    #[test]
    fn decrypt_message_with_header_and_base_64_encoding() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let message = b"3.133ed3d4a21effe90zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=";
        let expected =
            r#"{"devId":"002004265ccf7fb1b659","dps":{"1":false,"2":0},"t":1529442366,"s":8}"#
                .as_bytes()
                .to_owned();

        let decrypted = cipher.decrypt(message).unwrap();
        assert_eq!(&expected, &decrypted);
    }

    #[test]
    fn decrypt_message_with_version_threethree() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeThree);
        let message = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=".to_vec();
        let message = base64::engine::general_purpose::STANDARD
            .decode(message)
            .unwrap();
        let expected =
            r#"{"devId":"002004265ccf7fb1b659","dps":{"1":false,"2":0},"t":1529442366,"s":8}"#
                .as_bytes()
                .to_owned();

        let decrypted = cipher.decrypt(&message).unwrap();
        assert_eq!(&expected, &decrypted);
        // In the case of ThreeThree version,  the boolean it does not matter. It is always NOT
        // base64 encoded.
        let decrypted = cipher.decrypt(&message).unwrap();
        assert_eq!(&expected, &decrypted);
    }

    #[test]
    fn decrypt_message_without_header_and_base64_encoding() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let message = b"zrA8OK3r3JMiUXpXDWauNppY4Am2c8rZ6sb4Yf15MjM8n5ByDx+QWeCZtcrPqddxLrhm906bSKbQAFtT1uCp+zP5AxlqJf5d0Pp2OxyXyjg=";
        let expected =
            r#"{"devId":"002004265ccf7fb1b659","dps":{"1":false,"2":0},"t":1529442366,"s":8}"#
                .as_bytes()
                .to_owned();

        let decrypted = cipher.decrypt(message).unwrap();
        assert_eq!(&expected, &decrypted);
    }

    #[test]
    fn decrypt_message_where_payload_is_not_json() {
        let cipher = TuyaCipher::create(b"bbe88b3f4106d354", TuyaVersion::ThreeOne);
        let message = b"3.133ed3d4a21effe90rt1hJFzMJPF3x9UhPTCiXw==";
        let expected = "gw id invalid".as_bytes().to_owned();

        let decrypted = cipher.decrypt(message).unwrap();
        assert_eq!(&expected, &decrypted);
    }
}