Midea/src/discover.rs

use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
    str::from_utf8,
    time::Duration,
};

use anyhow::Result;

use crate::{hex, security::Security};

#[derive(Debug, Clone)]
pub struct DeviceInfo {
    pub id: u64,
    pub model: String,
    pub sn: String,
    pub protocol: u8,
    pub device_type: u8,

    pub addr: SocketAddr,
}

pub struct Startup;

impl Startup {
    const BROADCAST_MSG: &'static [u8] = &[
        0x5a, 0x5a, 0x01, 0x11, 0x48, 0x00, 0x92, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7f, 0x75, 0xbd, 0x6b, 0x3e,
        0x4f, 0x8b, 0x76, 0x2e, 0x84, 0x9c, 0x6e, 0x57, 0x8d, 0x65, 0x90, 0x03, 0x6e, 0x9d, 0x43,
        0x42, 0xa5, 0x0f, 0x1f, 0x56, 0x9e, 0xb8, 0xec, 0x91, 0x8e, 0x92, 0xe5,
    ];

    const NUM_RETRIES: u8 = 5;

    pub async fn discover() -> Result<Vec<DeviceInfo>> {
        let socket = UdpSocket::bind("0.0.0.0:0")?;
        socket.set_broadcast(true)?;
        socket.set_read_timeout(Some(Duration::from_secs(2)))?;

        let mut devices = HashMap::new();

        for addr in Self::broadcast_addresses() {
            for _ in 0..Self::NUM_RETRIES {
                socket.send_to(Self::BROADCAST_MSG, (addr, 6445))?;
                socket.send_to(Self::BROADCAST_MSG, (addr, 20086))?;

                let mut buffer = [0; 1024];
                if let Ok((bytes_read, mut addr)) = socket.recv_from(&mut buffer) {
                    let mut bytes = buffer[0..bytes_read].to_vec();

                    let protocol = if bytes[..2] == hex("5a5a")? {
                        2
                    } else if bytes[..2] == hex("8370")? {
                        if bytes[8..10] == hex("5a5a")? {
                            bytes = bytes[8..(bytes.len() - 16)].to_vec();
                        }

                        3
                    } else {
                        continue;
                    };

                    let mut tmp: Vec<u8> = bytes[20..26].to_vec();
                    tmp.push(0);
                    tmp.push(0);

                    let device_id = u64::from_le_bytes(tmp.try_into().unwrap());

                    if devices.contains_key(&device_id) {
                        continue;
                    }

                    let len = bytes.len();
                    let encrypt_data = &mut bytes[40..(len - 16)];
                    Security::aes_decrypt(encrypt_data);

                    let start = 41;
                    let upper = start + encrypt_data[40] as usize;
                    let ssid = from_utf8(&encrypt_data[start..upper])?;

                    let device_type = u8::from_str_radix(ssid.split('_').nth(1).unwrap(), 16)?;
                    let model = from_utf8(&encrypt_data[17..25])?;
                    let sn = from_utf8(&encrypt_data[8..40])?;
                    let port = Self::bytes_to_port(&encrypt_data[4..8]);

                    addr.set_port(port);

                    devices.insert(
                        device_id,
                        (
                            addr,
                            model.to_string(),
                            sn.to_string(),
                            protocol,
                            device_type,
                        ),
                    );
                }
            }
        }

        Ok(devices
            .into_iter()
            .map(
                |(id, (addr, model, sn, protocol, device_type))| DeviceInfo {
                    id,
                    model,
                    sn,
                    protocol,
                    device_type,

                    addr,
                },
            )
            .collect())
    }

    fn bytes_to_port(bytes: &[u8]) -> u16 {
        let mut b = 0;
        let mut i: u32 = 0;

        while b < 4 {
            let b1 = if b < bytes.len() {
                (bytes[b] & 0xFF) as u32
            } else {
                0
            };

            i |= b1 << b * 8;
            b += 1;
        }

        i as u16
    }
}

impl Startup {
    fn broadcast_addresses() -> Vec<IpAddr> {
        use if_addrs::*;

        let mut addresses = Vec::new();

        if let Some(interfaces) = get_if_addrs().ok() {
            for interface in interfaces {
                match interface.addr {
                    IfAddr::V4(v4) => {
                        let ip = v4.ip;
                        let octets = ip.octets();

                        if ip.is_private() {
                            addresses.push(IpAddr::V4(Ipv4Addr::new(
                                octets[0],
                                octets[1],
                                octets[2],
                                u8::MAX,
                            )))
                        }
                    }
                    IfAddr::V6(_v6) => {
                        // TODO
                        // addresses.push(IpAddr::V6(v6.ip));
                    }
                }
            }
        }

        addresses
    }
}

#[cfg(test)]
mod test {
    use anyhow::Result;

    use super::Startup;

    #[tokio::test]
    async fn discover() -> Result<()> {
        Startup::discover().await?;

        Ok(())
    }
}