1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use crate::store;
use diesel::prelude::*;
use keyring::Keyring;
use openssl::{
    hash,
    pkcs5,
    symm
};
use std::collections::HashMap;

/// LeetCode Cookies Schema
mod schema {
    table! {
        cookies (host_key) {
            encrypted_value -> Binary,
            host_key -> Text,
            name -> Text,
        }
    }
}

/// Please make sure the order
/// 
/// The order between table and struct must be same.
#[derive(Queryable, Debug, Clone)]
struct Cookies {
    pub encrypted_value: Vec<u8>,
    pub host_key: String,
    pub name: String,
}

/// Spawn cookies to cookie format
#[derive(Debug)]
pub struct Ident {
    pub csrf: String,
    session: String,
}

impl std::string::ToString for Ident {
    fn to_string(&self) -> String {
        format!(
            "LEETCODE_SESSION={};csrftoken={};",
            self.session,
            self.csrf
        ).to_string()
    }
}

/// Get cookies from chrome storage
pub fn cookies() -> Ident {
    use self::schema::cookies::dsl::*;
    debug!("Derive cookies from google chrome...");
    
    let home = dirs::home_dir().unwrap();
    let p = match std::env::consts::OS {
        "macos" => home.join("Library/Application Support/Google/Chrome/Default/Cookies"),
        "windows" => {
            let mut appd = std::path::PathBuf::new();
            let dir = app_dirs::get_data_root(app_dirs::AppDataType::SharedData);
            if dir.is_ok() {
                appd = dir.unwrap();
            }
            
            appd.join("../Local/Google/Chrome/User Data/Default/Cookies")
        },
        _ => home.join(".config/google-chrome/Default/Cookies"),
    };
    
    let conn = store::conn(p.to_string_lossy().to_string());
    let res = cookies
        .filter(host_key.like("%leetcode.com"))
        .load::<Cookies>(&conn)
        .expect("Loading cookies from google chrome failed.");

    // Get system password
    let ring = Keyring::new("Chrome Safe Storage", "Chrome");
    let pass = ring.get_password().expect("Get Password failed");

    // Decode cookies
    let mut m: HashMap<String, String> = HashMap::new();
    for c in res.to_owned() {
        if (
            c.name == "csrftoken".to_string()
        ) || (
            c.name == "LEETCODE_SESSION".to_string()
        ) {
            m.insert(c.name, decode_cookies(&pass, c.encrypted_value));
        }
    }
    
    Ident {
        csrf: m.get("csrftoken").unwrap().to_string(),
        session: m.get("LEETCODE_SESSION").unwrap().to_string(),
    }
}


/// Decode cookies from chrome
fn decode_cookies(pass: &str, v: Vec<u8>) -> String {
    let mut key = [0_u8; 16];
    pkcs5::pbkdf2_hmac(
        pass.as_bytes(),
        b"saltysalt",
        1003,
        hash::MessageDigest::sha1(),
        &mut key
    ).expect("pbkdf2 hmac went error.");
    
    chrome_decrypt(v, key)
}


/// Decrypt chrome cookie value with aes-128-cbc
fn chrome_decrypt(v: Vec<u8>, key: [u8;16]) -> String {
    // <space>: \u16
    let iv = vec![32_u8; 16];
    let mut decrypter = symm::Crypter::new(
        symm::Cipher::aes_128_cbc(),
        symm::Mode::Decrypt,
        &key,
        Some(&iv)
    ).unwrap();

    
    let data_len = v.len() - 3;
    let block_size = symm::Cipher::aes_128_cbc().block_size();
    let mut plaintext = vec![0; data_len + block_size];
    
    decrypter.pad(false);

    let count = decrypter.update(&v[3..], &mut plaintext).unwrap();
    decrypter.finalize(&mut plaintext[count..]).unwrap();
    plaintext.retain(|x| x >= &20_u8);

    String::from_utf8_lossy(&plaintext.to_vec()).to_string()
}