dlqrcode/src/lib.rs

extern crate base64_url;
extern crate hex;

use std::fmt;
use std::slice;
use std::time::Instant;
use std::time::SystemTime;

use sha2::{Digest, Sha256};

use aes::Aes256;
use block_modes::block_padding::Pkcs7;
use block_modes::{BlockMode, Cbc};
use totp_rs::{Algorithm, TOTP};

type Aes256Cbc = Cbc<Aes256, Pkcs7>;

const TOTP_STEP: u64 = 5;
const TOTP_SKEW: u8 = 3;
static QR_FIELD_DILIMITER: &str = ":";

pub struct DaliQrCode {
    master_key: Vec<u8>,
    iv: Vec<u8>,
    totp_step: u64,
    totp_skew: u8,
    totp_seed: Vec<u8>,
}

#[derive(Debug)]
pub struct DaliQrData {
    pub uid: String,
    pub cardno: String,
    pub cardtype: String,
    pub totp: String,
    pub nonce: String,
    sign: Vec<u8>,
}

impl DaliQrData {
    #[allow(dead_code)]
    fn new() -> Self {
        Self {
            uid: Default::default(),
            cardno: Default::default(),
            cardtype: Default::default(),
            totp: Default::default(),
            nonce: Default::default(),
            sign: Vec::new(),
        }
    }

    fn from_qrcode(qr_fields: &Vec<Vec<u8>>) -> Result<Self> {
        if qr_fields.len() < 6 {
            return Err(DecodeError::new("qrcode fields length must grater than 6."));
        }
        let sign = qr_fields[5].to_vec();
        Ok(Self {
            uid: String::from_utf8_lossy(&qr_fields[0].as_slice()).to_string(),
            cardno: String::from_utf8_lossy(&qr_fields[1].as_slice()).to_string(),
            cardtype: String::from_utf8_lossy(&qr_fields[2].as_slice()).to_string(),
            totp: String::from_utf8_lossy(&qr_fields[3].as_slice()).to_string(),
            nonce: String::from_utf8_lossy(&qr_fields[4].as_slice()).to_string(),
            sign: sign,
        })
    }

    fn update_sign(&mut self, sign: &Vec<u8>) {
        self.sign = sign.to_vec();
    }

    fn to_qrdata(&self) -> String {
        let v = vec![
            String::from(&self.uid),
            String::from(&self.cardno),
            String::from(&self.cardtype),
            String::from(&self.totp),
            String::from(&self.nonce),
        ];
        v.join(QR_FIELD_DILIMITER)
    }
}

#[derive(Debug, Clone)]
pub struct DecodeError {
    message: String,
}

impl fmt::Display for DecodeError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Decode Qrcode error {}", self.message)
    }
}

impl DecodeError {
    fn new(message: &str) -> Self {
        Self {
            message: String::from(message),
        }
    }
}

type Result<T> = std::result::Result<T, DecodeError>;

const KEY_LEN: usize = 32;
impl DaliQrCode {
    pub fn new(
        key: [u8; KEY_LEN],
        iv: Option<[u8; 16]>,
        step: Option<u64>,
        skew: Option<u8>,
        seed: Option<Vec<u8>>,
    ) -> Result<Self> {
        let key = key.to_vec();
        if key.len() != KEY_LEN {
            return Err(DecodeError::new(&format!(
                "key size must be {} bytes",
                KEY_LEN
            )));
        }
        let iv = if let Some(v) = iv {
            v.to_vec()
        } else {
            hex::decode("55b6f5b3287c535f8274b99354676d0e").unwrap()
        };

        if iv.len() != 16 {
            return Err(DecodeError::new("IV size must be 16 bytes"));
        }

        let totp_step = step.unwrap_or(TOTP_STEP);
        let totp_skew = skew.unwrap_or(TOTP_SKEW);
        let totp_seed = seed.unwrap_or(hex::decode("125ea2f97689988b6501").unwrap());

        Ok(Self {
            master_key: key,
            iv,
            totp_seed,
            totp_skew,
            totp_step,
        })
    }

    fn split_qrdata(&self, qrcode: &Vec<u8>, qrdata: &mut Vec<Vec<u8>>) -> Result<()> {
        let v = QR_FIELD_DILIMITER.as_bytes()[0];
        let mut j = 0;
        for i in 0..qrcode.len() {
            if qrcode[i] == v {
                if i > j {
                    let mut s = Vec::new();
                    s.extend_from_slice(&qrcode[j..i]);
                    qrdata.push(s);
                }
                j = i + 1;
            }
        }
        if j < qrcode.len() {
            let mut s = Vec::new();
            s.extend_from_slice(&qrcode[j..]);
            qrdata.push(s);
        }
        Ok(())
    }

    pub fn decode(&self, qrcode: *const u8, len: usize, secs_offset: i32) -> Result<DaliQrData> {
        let qrcode = self.decode_qrcode(qrcode, len)?;
        let mut qr_fields: Vec<Vec<u8>> = Vec::new();

        self.split_qrdata(&qrcode, &mut qr_fields)?;
        
        let qr_data = DaliQrData::from_qrcode(&qr_fields)?;
        if qr_data.uid.len() < 16 {
            return Err(DecodeError::new("uid must grater than 16"));
        }

        self.check_qrcode_sign(&qr_data)?;

        let totp = self.new_totp();
        let time = self.totp_time(secs_offset);
        if totp.check(&qr_data.totp, time) {
            Ok(qr_data)
        } else {
            Err(DecodeError::new("qrcode totp error"))
        }
    }

    fn check_qrcode_sign(&self, qr_data: &DaliQrData) -> Result<()> {
        let sign = self.calc_sign(qr_data);
        if qr_data.sign == sign {
            Ok(())
        } else {
            Err(DecodeError::new("sign is invalid"))
        }
    }

    fn decode_qrcode(
        &self,
        qrcode: *const u8,
        len: usize,
    ) -> Result<Vec<u8>> {
        let cipher = match Aes256Cbc::new_var(&self.master_key, &self.iv) {
            Ok(c) => c,
            Err(e) => return Err(DecodeError::new(&format!("aes key error {:?}", e))),
        };

        let qrcode = unsafe {
            let s = slice::from_raw_parts(qrcode, len);
            if let Ok(code) = base64_url::decode(s) {
                code
            } else {
                return Err(DecodeError::new("data base64 decode error"));
            }
        };

        match cipher.decrypt_vec(&qrcode) {
            Ok(data) => Ok(data),
            Err(e) => Err(DecodeError::new(&format!("block error {:?}", e))),
        }
    }

    fn totp_time(&self, secs_offset: i32) -> u64 {
        let time = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        if secs_offset > 0 {
            time + (secs_offset as u64)
        } else {
            time - (secs_offset as u64)
        }
    }

    fn new_totp(&self) -> TOTP<Vec<u8>> {
        let seed = self.totp_seed.clone();
        TOTP::new(Algorithm::SHA1, 8, self.totp_skew, self.totp_step, seed)
    }

    fn encode_qrcode(&self, qr_data: &DaliQrData) -> Result<String> {
        let plain_text = qr_data.to_qrdata();
        let cipher = match Aes256Cbc::new_var(&self.master_key, &self.iv) {
            Ok(c) => c,
            Err(e) => return Err(DecodeError::new(&format!("aes key error {:?}", e))),
        };
        let mut buffer = Vec::new();
        buffer.extend_from_slice(&plain_text.as_bytes());
        buffer.push(QR_FIELD_DILIMITER.as_bytes()[0]);
        buffer.extend_from_slice(qr_data.sign.as_slice());
        let crypt_data = cipher.encrypt_vec(buffer.as_slice());
        Ok(base64_url::encode(crypt_data.as_slice()))
    }

    fn calc_sign(&self, qr_data: &DaliQrData) -> Vec<u8> {
        let mut hasher = Sha256::new();
        hasher.update("{dlsmk_}".as_bytes());
        hasher.update(qr_data.uid.as_bytes());
        let salt = hasher.finalize();
        let mut hasher = Sha256::new();
        hasher.update(qr_data.to_qrdata().as_bytes());
        hasher.update(salt);
        hasher.finalize().to_vec()
    }

    pub fn encode(&self, qr_data: &mut DaliQrData, secs_offset: i32) -> Result<String> {
        if qr_data.nonce.len() == 0 {
            qr_data.nonce = format!("{:02}", Instant::now().elapsed().as_secs() % 100);
        }
        if qr_data.totp.len() == 0 {
            let totp = self.new_totp();
            let time = self.totp_time(secs_offset);
            qr_data.totp = totp.generate(time);
        }
        let sign = self.calc_sign(qr_data);
        qr_data.update_sign(&sign);
        self.encode_qrcode(qr_data)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::decode;
    use std::convert::TryInto;
    const KEYLEN: usize = 32;
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }

    #[test]
    fn aes_test() {
        let mut key = [0u8; KEYLEN];
        let s = decode("Vbb1syh8U1+CdLmTVGdtDiVvKBQ81n4GmgBEO/ohSbU=").unwrap();
        key.clone_from_slice(&s.as_slice()[..KEYLEN]);

        let iv: [u8; 16] = {
            let s = hex::decode("55b6f5b3287c535f8274b99354676d0e").unwrap();
            s.into_boxed_slice().as_ref().try_into().unwrap()
        };

        let aes_key = Aes256Cbc::new_var(&key, &iv).unwrap();
        let plaintext = String::from("hello ldldldf ldfldl dslfasdamf sdmfdfdf");

        let mut buffer = Vec::new();
        buffer.extend_from_slice(plaintext.as_bytes());
        println!(
            "plain len : {} , buffer len : {}",
            plaintext.len(),
            buffer.len()
        );
        let cipher_data = aes_key.encrypt_vec(&buffer);

        println!(
            "cipher len : {}, last {}",
            cipher_data.len(),
            cipher_data[cipher_data.len() - 1] as u32
        );

        let aes_key = Aes256Cbc::new_var(&key, &iv).unwrap();
        let _ = aes_key.decrypt_vec(&cipher_data);
    }

    fn get_key() -> ([u8; KEYLEN], [u8; 16]) {
        let mut key = [0u8; KEYLEN];
        let s = decode("Vbb1syh8U1+CdLmTVGdtDiVvKBQ81n4GmgBEO/ohSbU=").unwrap();
        key.clone_from_slice(&s.as_slice()[..KEYLEN]);

        let iv: [u8; 16] = {
            let s = hex::decode("55b6f5b3287c535f8274b99354676d0e").unwrap();
            s.into_boxed_slice().as_ref().try_into().unwrap()
        };
        (key, iv)
    }

    #[test]
    fn check_qrcode_encode() {
        let expect_qrcode = String::from("6lHyFX_vg5U2hymn8OsdNUD7dT0-sCmEQkKrm9cnzHlku6-FYxuL6nP5YR2Fve8Sfj-Asd-3dfQUkaiqqbfQWO8B_811B3uhHmGm9IjlpLicz_c1H1_ORb9tJl-IhMKu");
        // let buffer = base64_url::decode(expect_qrcode.as_bytes()).unwrap();

        // println!("encrypt buffer <{}>", buffer.len());
        // println!("decode b64<{}>", hex::encode(buffer.clone()));
        // let aes_key = Aes256Cbc::new_var(&key, &iv).unwrap();
        // let data = aes_key.decrypt_vec(&buffer).unwrap();
        // println!("data : {}", String::from_utf8_lossy(data.as_slice()));

        let (key, iv) = get_key();

        let mut qr_data = DaliQrData::new();
        qr_data.uid = String::from("0a5de6ce985d43989b7ebe64ad8eb9c3");
        qr_data.cardno = String::from("00001252");
        qr_data.cardtype = String::from("80");
        qr_data.nonce = String::from("ac");
        qr_data.totp = String::from("50053019");

        let dali_qrcode = DaliQrCode::new(key, Some(iv), None, None, None).unwrap();

        match dali_qrcode.encode(&mut qr_data, 0) {
            Ok(qrcode) => {
                assert_eq!(qrcode, expect_qrcode);
            }
            Err(e) => {
                panic!("error {}", e);
            }
        }
    }

    #[test]
    fn check_qrcode_decode() {
        let (key, iv) = get_key();
        let mut qr_data = DaliQrData::new();
        qr_data.uid = String::from("0a5de6ce985d43989b7ebe64ad8eb9c3");
        qr_data.cardno = String::from("00001252");
        qr_data.cardtype = String::from("80");

        let dali_qrcode = DaliQrCode::new(key, Some(iv), Some(60u64), Some(5u8), None).unwrap();

        match dali_qrcode.encode(&mut qr_data, 0) {
            Ok(qrcode) => {
                if let Err(e) = dali_qrcode.decode(qrcode.as_ptr(), qrcode.len(), 0) {
                    panic!("error {}", e);
                }
            }
            Err(e) => {
                panic!("error {}", e);
            }
        }
    }
}