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
137
138
139
140
141
142
143
144
145
146
147
148

use crate::cache;
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)
    }
}

/// Get cookies from chrome storage
pub fn cookies() -> Result<Ident, crate::Error> {
    let ccfg = crate::cfg::locate()?.cookies;
    if !ccfg.csrf.is_empty() && !ccfg.session.is_empty() {
        return Ok(Ident {
            csrf: ccfg.csrf,
            session: ccfg.session,
        });
    }

    // If doesn't config SESSION and csrftoken
    use self::schema::cookies::dsl::*;
    trace!("Derive cookies from google chrome...");

    let home = dirs::home_dir()?;
    let p = match std::env::consts::OS {
        "macos" => home.join("Library/Application Support/Google/Chrome/Default/Cookies"),
        "linux" => home.join(".config/google-chrome/Default/Cookies"),
        _ => panic!("Opps...only works on OSX or Linux now..."),
    };

    debug!("Chrome Cookies path is {:?}", &p);
    let conn = cache::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.");

    debug!("res {:?}", &res);
    if res.len() == (0 as usize) {
        return Err(crate::Error::CookieError);
    }

    // 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") || (c.name == "LEETCODE_SESSION") {
            m.insert(c.name, decode_cookies(&pass, c.encrypted_value)?);
        }
    }

    Ok(Ident {
        csrf: m.get("csrftoken")?.to_string(),
        session: m.get("LEETCODE_SESSION")?.to_string(),
    })
}

/// Decode cookies from chrome
fn decode_cookies(pass: &str, v: Vec<u8>) -> Result<String, crate::Error> {
    let mut key = [0_u8; 16];
    match std::env::consts::OS {
        "macos" => {
            pkcs5::pbkdf2_hmac(
                pass.as_bytes(),
                b"saltysalt",
                1003,
                hash::MessageDigest::sha1(),
                &mut key,
            )
            .expect("pbkdf2 hmac went error.");
        }
        "linux" => {
            pkcs5::pbkdf2_hmac(
                b"peanuts",
                b"saltysalt",
                1,
                hash::MessageDigest::sha1(),
                &mut key,
            )
            .expect("pbkdf2 hmac went error.");
        }
        _ => {
            return Err(crate::Error::FeatureError(
                "only supports OSX or Linux for now".to_string(),
            ))
        }
    }

    chrome_decrypt(v, key)
}

/// Decrypt chrome cookie value with aes-128-cbc
fn chrome_decrypt(v: Vec<u8>, key: [u8; 16]) -> Result<String, crate::Error> {
    // <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),
    )?;

    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)?;
    decrypter.finalize(&mut plaintext[count..])?;
    plaintext.retain(|x| x >= &20_u8);

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