WBI Sign Implementation

playground.rs

// Last updated: 2023/12/25 22:56
// Gist: https://gist.github.com/cxw620/81bbb3b59fd32255764161775ab97097

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

    static IMG_KEY: &'static str = "7cd084941338484aae1ad9425b84077c";
    static SUB_KEY: &'static str = "4932caff0ff746eab6f01bf08b70ac45";
    // static IMG_KEY: &'static str = "4a1d4479a1ea4146bc7552eea71c28e9";
    // static SUB_KEY: &'static str = "fa5812e23a204d10b332dc24d992432d";
    #[test]
    fn test_wbi_sign() {
        // let query_vec = vec![
        //     ("__refresh__", "true".into()),
        //     ("_extra", "".into()),
        //     ("context", "".into()),
        //     ("page", "1".into()),
        //     ("page_size", "42".into()),
        //     ("order", "".into()),
        //     ("duration", "".into()),
        //     ("from_source", "".into()),
        //     ("from_spmid", "333.337".into()),
        //     ("platform", "pc".into()),
        //     ("highlight", "1".into()),
        //     ("single_column", "0".into()),
        //     ("keyword", "克洛琳德".into()),
        //     ("qv_id", "iJNVBR35jvcl2pRehSOaRimrin1Ufn85".into()),
        //     ("ad_resource", "5646".into()),
        //     ("source_tag", "3".into()),
        //     ("web_location", "1430654".into()),
        //     // ("gaia_vtoken", "".into()),
        // ];
        // let query_vec = vec![
        //     ("foo", "114".into()),
        //     ("bar", "514".into()),
        //     ("zab", "1919810".into()),
        // ];
        let query_vec = vec![
            ("mid", "11997177".into()),
            ("token", "".into()),
            ("platform", "web".into()),
            ("web_location", "1550101".into()),
        ];
        let w_rid = Wbi {
            query_vec,
            img_key: IMG_KEY,
            sub_key: SUB_KEY,
        }
        .gen_w_rid(1703513649)
        .unwrap();
        assert_eq!(w_rid, "7d4428b3f2f9ee2811e116ec6fd41a4f");
    }
}

use std::borrow::Cow;

/// WBI Sign implementation V1.0.1
///
/// Last updated: 23-12-05 04:09
pub struct Wbi<'q> {
    img_key: &'q str,
    sub_key: &'q str,
    query_vec: Vec<(&'q str, Cow<'q, str>)>,
}

#[allow(dead_code)]
impl<'q> Wbi<'q> {
    pub fn new(
        img_key: &'q str,
        sub_key: &'q str,
        query_vec: Vec<(&'q str, Cow<'q, str>)>,
    ) -> Self {
        Self {
            img_key,
            sub_key,
            query_vec,
        }
    }

    /// Returns signed url query string.
    pub fn sign(mut self, wts: Option<u64>) -> Result<String, BoxedError> {
        let wbi_key = str_concat!(self.img_key, self.sub_key);
        if wbi_key.len() != 64 {
            return Err(SignErr::InvalidWbiKey.into());
        }

        let mixin_key = Self::gen_mixin_key(wbi_key);

        let wts: String = wts.map_or_else(|| now!().as_secs(), |wts| wts).to_string();
        let wts_param = str_concat!("wts=", &wts);

        self.query_vec.push(("wts", wts.into()));

        let unsigned_query = QueryBuilder::encode(self.query_vec, true);

        let w_rid = calc_md5!(str_concat!(&unsigned_query, &mixin_key));

        let signed_query = {
            // `wts_param` will and will only appear once in unsigned_query
            let (mut start, part) = unsigned_query.match_indices(&wts_param).next().unwrap();
            let mut part_len = part.len();
            // `start` > 0 then not the first, should also remove `&` before `wts`
            if start != 0 {
                start -= 1;
                part_len += 1;
            }
            str_concat!(
                unsafe { unsigned_query.get_unchecked(0..start) },
                unsafe { unsigned_query.get_unchecked((start + part_len)..unsigned_query.len()) },
                "&w_rid=",
                &w_rid,
                "&",
                &wts_param
            )
        };

        Ok(signed_query)
    }

    #[cfg(test)]
    fn gen_w_rid(mut self, wts: impl ToString) -> Result<String, BoxedError> {
        let wbi_key = str_concat!(self.img_key, self.sub_key);
        let mixin_key = Self::gen_mixin_key(wbi_key);
        println!("mixin_key: {}", mixin_key);
        let wts: String = wts.to_string();
        self.query_vec.push(("wts", wts.into()));
        self.query_vec.sort_unstable_by_key(|param| param.0);
        let unsigned_query = QueryBuilder::encode(self.query_vec, true);
        let w_rid = calc_md5!(str_concat!(&unsigned_query, &mixin_key));
        Ok(w_rid)
    }

    /// Gen mixin key
    #[inline]
    fn gen_mixin_key(raw_wbi_key: impl AsRef<[u8]>) -> String {
        const MIXIN_KEY_ENC_TAB: [u8; 64] = [
            46, 47, 18, 2, 53, 8, 23, 32, 15, 50, 10, 31, 58, 3, 45, 35, 27, 43, 5, 49, 33, 9, 42,
            19, 29, 28, 14, 39, 12, 38, 41, 13, 37, 48, 7, 16, 24, 55, 40, 61, 26, 17, 0, 1, 60,
            51, 30, 4, 22, 25, 54, 21, 56, 59, 6, 63, 57, 62, 11, 36, 20, 34, 44, 52,
        ];
        let raw_wbi_key = raw_wbi_key.as_ref();
        let mut mixin_key_raw = {
            let binding = MIXIN_KEY_ENC_TAB
                .iter()
                .map(|n| raw_wbi_key[*n as usize])
                .collect::<Vec<u8>>();
            unsafe { String::from_utf8_unchecked(binding) }
        };
        let _ = mixin_key_raw.split_off(32);
        mixin_key_raw
    }

    /// For future use.
    fn swap_string(input: &str, t: u32) -> String {
        if input.len() % 2 != 0 {
            return input.to_owned();
        }
        if t == 0 {
            return input.to_owned();
        }
        if input.len() == 2u32.pow(t) as usize {
            return input.chars().rev().collect();
        }
        let mid = input.len() / 2;
        let r = &input[..mid];
        let n = &input[mid..];
        str_concat!(&Self::swap_string(n, t - 1), &Self::swap_string(r, t - 1))
    }
}

// -------------------------------------------------------------------------------
// The following are copied from my private project to make this example compiled.
// You can just ignore them.
// -------------------------------------------------------------------------------

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Signer<'s> {
    App {
        appkey: &'s str,
        appsec: Option<&'s str>,
    },
    Wbi {
        img_key: &'s str,
        sub_key: &'s str,
    },
    None,
}

pub struct QueryBuilder<'q> {
    pub query_vec: Vec<(&'q str, Cow<'q, str>)>,
    pub signer: Signer<'q>,
    pub need_sort: bool,
}

#[allow(unused_variables)]
impl<'q> QueryBuilder<'q> {
    pub fn build(self) -> Result<String, BoxedError> {
        match self.signer {
            Signer::App { appkey, appsec } => {
                // sign::App::new(appkey, appsec, self.query_vec).sign(Some(now!().as_secs()))
                unimplemented!("APP Sign not implmented in this example.")
            }
            Signer::Wbi { img_key, sub_key } => {
                // sign::Wbi::new(img_key, sub_key, self.query_vec).sign(None)
                Wbi::new(img_key, sub_key, self.query_vec).sign(None)
            }
            Signer::None => Ok(Self::encode(self.query_vec, self.need_sort)),
        }
    }
    /// Encode query from given vec.
    ///
    /// # WARNING
    ///
    /// For performance:
    /// - **KEY will not be url encoded**, for all query key is known and without special chars.
    /// - When `need_sort`, the given vec will be sorted in place. **Do clone a new one in advance**
    ///   if you need the original one.
    // #[inline]
    // pub fn encode(mut query_vec: Vec<(&'q str, Cow<'q, str>)>, need_sort: bool) -> String {
    //     if need_sort {
    //         Self::sort(&mut query_vec);
    //     }
    //     let mut query_string = String::with_capacity(256);
    //     query_vec.into_iter().for_each(|(k, v)| {
    //         query_string.push_str(k);
    //         query_string.push('=');
    //         query_string.push_str(&urlencoding::encode(v.as_ref()));
    //         query_string.push('&');
    //     });
    //     query_string.pop();
    //     query_string
    // }
    // #[inline]
    // fn sort(query_vec: &mut Vec<(&str, Cow<'_, str>)>) {
    //     query_vec.sort_unstable_by_key(|param| param.0);
    // }
    #[inline]
    pub fn encode(query_vec: Vec<(&'q str, Cow<'q, str>)>, need_sort: bool) -> String {
        // ! Space will be encoded as `+` instead of `%20` in crate `serde_urlencoded`.
        serde_urlencoded::to_string(query_vec).unwrap().replace("+", "%20")
    }
}

// Error handler

pub type BoxedError = Box<dyn std::error::Error + Send + Sync>;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("{0}")]
    Sign(
        #[from]
        #[source]
        SignErr,
    ),
}

#[derive(Debug, thiserror::Error)]
pub enum SignErr {
    #[error("Invalid given wbi key, length not 64")]
    InvalidWbiKey,
    #[error("Unknown appkey {0}, corresponding appsec not found")]
    UnknownAppkey(String),
}


// Some macros with specific function, can be replaced with your own one

#[macro_export]
macro_rules! str_concat {
    ($($x:expr),*) => {
        {
            let mut string_final = String::with_capacity(512);
            $(
                string_final.push_str($x);
            )*
            string_final
        }
    };
}

#[macro_export]
macro_rules! calc_md5 {
    ($input_str: expr) => {{
        // let mut md5_instance = crypto::md5::Md5::new();
        // crypto::digest::Digest::input_str(&mut md5_instance, &($input_str));
        // crypto::digest::Digest::result_str(&mut md5_instance)
        use md5::{Digest, Md5};
        let mut hasher = Md5::new();
        hasher.update(&($input_str));
        let result = hasher.finalize();
        format!("{:0>2x}", result)
    }};
}

#[macro_export]
/// Faster way to get current timestamp other than `chrono::Local::now().timestamp()`,
/// 12x faster on my machine.
///
/// See [`Duration`](https://doc.rust-lang.org/std/time/struct.Duration.html) for more details.
macro_rules! now {
    () => {{
        match std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
            Ok(t) => t,
            Err(_) => panic!("SystemTime before UNIX EPOCH!"),
        }
    }};
}