Midea/src/security.rs

use aes::{
    cipher::{
        block_padding::Pkcs7, generic_array::GenericArray, BlockDecrypt, BlockDecryptMut,
        BlockEncryptMut, KeyInit, KeyIvInit,
    },
    Aes128,
};
use anyhow::{bail, Result};
use rand::{self, RngCore};
use sha2::{Digest, Sha256};

use crate::hex;

#[allow(non_camel_case_types)]
#[repr(u8)]
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
pub enum MsgType {
    HANDSHAKE_REQUEST = 0x0,
    ANDSHAKE_RESPONSE = 0x1,
    ENCRYPTED_RESPONSE = 0x3,
    ENCRYPTED_REQUEST = 0x6,
}

#[derive(Debug, Default)]
pub struct Security {
    request_count: u16,
    response_count: u16,

    tcp_key: Option<[u8; 32]>,
}

impl Security {
    const N: u128 = 141661095494369103254425781617665632877;
    const KEY: [u8; 16] = Self::N.to_be_bytes();
    const IV: [u8; 16] = [b'\0'; 16];

    pub fn aes_decrypt(data: &mut [u8]) -> &[u8] {
        let array = GenericArray::from(Self::KEY);
        let cipher = Aes128::new(&array);

        for chunk in data.chunks_mut(16) {
            let mut block = GenericArray::from_mut_slice(chunk);
            cipher.decrypt_block(&mut block);
        }

        data
    }

    pub fn aes_cbc_encrypt(&self, raw: &[u8], key: &[u8; 32]) -> Vec<u8> {
        type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;

        let mut buf = vec![0; raw.len()];
        buf.copy_from_slice(raw);

        Aes256CbcEnc::new(key.into(), &Self::IV.into())
            .encrypt_padded_mut::<Pkcs7>(&mut buf, raw.len())
            .unwrap();

        buf
    }

    pub fn aes_cbc_decrypt(&self, raw: &[u8], key: &[u8; 32]) -> Vec<u8> {
        type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

        let mut buf = vec![0; raw.len()];
        buf.copy_from_slice(raw);

        Aes256CbcDec::new(key.into(), &Self::IV.into())
            .decrypt_padded_mut::<Pkcs7>(&mut buf)
            .unwrap();

        buf
    }

    pub fn tcp_key(&mut self, response: &[u8], key: &[u8; 32]) -> Result<()> {
        if response == b"ERROR" {
            bail!("authentication failed!");
        }

        let payload = &response[0..32];
        let sign = &response[32..];

        let plain = self.aes_cbc_decrypt(payload, &key);

        if Sha256::digest(&plain).into_iter().collect::<Vec<u8>>() != sign {
            bail!("sign does not match");
        }

        self.tcp_key = Some(Self::xorstr(&plain, key).try_into().unwrap());
        self.request_count = 0;
        self.response_count = 0;

        Ok(())
    }

    fn xorstr(lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
        assert_eq!(lhs.len(), rhs.len());

        lhs.iter().zip(rhs.iter()).map(|(&l, &r)| l ^ r).collect()
    }

    pub fn encode_8370(&mut self, msg: &[u8], msg_type: MsgType) -> Result<Vec<u8>> {
        let mut header = hex("8370")?;
        let mut data: Vec<u8> = msg.to_vec();

        let mut size = data.len() as u16;
        let mut padding = 0;

        if msg_type == MsgType::ENCRYPTED_RESPONSE || msg_type == MsgType::ENCRYPTED_REQUEST {
            if (size + 2) % 16 != 0 {
                padding = 16 - ((size + 2) & 0xf);
                size += padding + 32;
                data.extend({
                    let mut d = vec![0; padding as usize];
                    rand::thread_rng().fill_bytes(&mut d);

                    d
                });
            }
        }

        header.extend(size.to_be_bytes());
        header.extend([0x20, (padding << 4) as u8 | msg_type as u8]);

        data = {
            let mut b = self.request_count.to_be_bytes().to_vec();
            b.extend(data);
            b
        };

        (self.request_count, _) = self.request_count.overflowing_add(1);

        if msg_type == MsgType::ENCRYPTED_RESPONSE || msg_type == MsgType::ENCRYPTED_REQUEST {
            let sign: Vec<u8> = Sha256::digest(Self::add_bytes(header.clone(), data.clone()))
                .into_iter()
                .collect();

            data = self.aes_cbc_encrypt(&data, self.tcp_key.as_ref().unwrap());
            data.extend(sign);
        }

        header.extend(data);
        Ok(header)
    }

    fn add_bytes(mut lhs: Vec<u8>, rhs: Vec<u8>) -> Vec<u8> {
        lhs.extend(rhs);
        lhs
    }
}