kleis 0.1.0

use std::collections::HashSet;
use std::env;
use std::fs;
use std::path::PathBuf;
use std::process::{self, Command};

use aes::Aes128;
use cbc::cipher::{block_padding::NoPadding, BlockDecryptMut, KeyIvInit};
use pbkdf2::pbkdf2_hmac;
use rusqlite::Connection;
use serde::Serialize;
use sha1::Sha1;

type Aes128CbcDec = cbc::Decryptor<Aes128>;

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Cookie {
    name: String,
    value: String,
    domain: String,
    path: String,
    secure: bool,
    http_only: bool,
}

fn chrome_cookies_path() -> PathBuf {
    let home = env::var("HOME").expect("HOME not set");
    PathBuf::from(home).join("Library/Application Support/Google/Chrome/Default/Cookies")
}

fn get_chrome_key() -> [u8; 16] {
    let output = Command::new("security")
        .args(["find-generic-password", "-s", "Chrome Safe Storage", "-w"])
        .output()
        .expect("failed to run `security` command");

    if !output.status.success() {
        eprintln!("Failed to get Chrome Safe Storage key from keychain.");
        eprintln!("Is your keychain unlocked? Try: security unlock-keychain");
        process::exit(1);
    }

    let password = String::from_utf8(output.stdout)
        .expect("non-UTF-8 keychain output")
        .trim()
        .to_string();

    let mut key = [0u8; 16];
    pbkdf2_hmac::<Sha1>(password.as_bytes(), b"saltysalt", 1003, &mut key);
    key
}

fn pkcs7_unpad(data: &[u8]) -> &[u8] {
    if data.is_empty() {
        return data;
    }
    let pad_len = *data.last().unwrap() as usize;
    if pad_len == 0 || pad_len > 16 || pad_len > data.len() {
        return data;
    }
    if data[data.len() - pad_len..]
        .iter()
        .all(|&b| b as usize == pad_len)
    {
        &data[..data.len() - pad_len]
    } else {
        data
    }
}

fn decrypt_v10(encrypted: &[u8], key: &[u8; 16]) -> Option<String> {
    if encrypted.len() <= 3 {
        return None;
    }
    if &encrypted[..3] != b"v10" {
        return Some(String::from_utf8_lossy(encrypted).to_string());
    }

    let ciphertext = &encrypted[3..];
    if ciphertext.is_empty() || ciphertext.len() % 16 != 0 {
        return None;
    }

    let iv = [b' '; 16];
    let mut buf = ciphertext.to_vec();

    let decryptor = Aes128CbcDec::new(key.into(), &iv.into());
    let decrypted = decryptor
        .decrypt_padded_mut::<NoPadding>(&mut buf)
        .ok()?;

    let decrypted = pkcs7_unpad(decrypted);

    // First AES block(s) may decrypt to garbage on newer Chrome.
    // Find longest trailing ASCII suffix — that's the real value.
    let mut best_start = decrypted.len();
    for i in (0..decrypted.len()).rev() {
        if decrypted[i] >= 0x20 && decrypted[i] < 0x7F {
            best_start = i;
        } else {
            break;
        }
    }

    let value = if best_start < decrypted.len() {
        String::from_utf8_lossy(&decrypted[best_start..]).to_string()
    } else {
        String::from_utf8_lossy(decrypted).to_string()
    };

    if value.is_empty() {
        None
    } else {
        Some(value)
    }
}

fn expand_domains(args: &[String]) -> Vec<String> {
    let mut domains = HashSet::new();
    for arg in args {
        domains.insert(arg.clone());
        if !arg.starts_with('.') {
            domains.insert(format!(".{arg}"));
        }
        let bare = arg.trim_start_matches('.');
        if !bare.starts_with("www.") {
            domains.insert(format!("www.{bare}"));
        }
    }
    domains.into_iter().collect()
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();

    if args.is_empty() || args[0] == "-h" || args[0] == "--help" {
        eprintln!("kleis — extract cookies from Chrome on macOS");
        eprintln!();
        eprintln!("Usage: kleis <domain> [domain2 ...]");
        eprintln!("Example: kleis midjourney.com");
        eprintln!("         kleis .example.com www.example.com");
        process::exit(if args.is_empty() { 1 } else { 0 });
    }

    let domains = expand_domains(&args);
    let key = get_chrome_key();

    // Copy cookie DB to temp file (Chrome may have it locked)
    let db_path = chrome_cookies_path();
    if !db_path.exists() {
        eprintln!("Chrome cookie database not found at {}", db_path.display());
        eprintln!("Is Chrome installed?");
        process::exit(1);
    }

    let tmp = tempfile::NamedTempFile::new().expect("failed to create temp file");
    fs::copy(&db_path, tmp.path()).expect("failed to copy cookie database");

    let conn = Connection::open(tmp.path()).expect("failed to open cookie database");
    let placeholders: String = domains.iter().map(|_| "?").collect::<Vec<_>>().join(",");
    let query = format!(
        "SELECT name, encrypted_value, host_key, path, is_secure, is_httponly \
         FROM cookies WHERE host_key IN ({placeholders})"
    );

    let mut stmt = conn.prepare(&query).expect("failed to prepare query");
    let params: Vec<&dyn rusqlite::types::ToSql> = domains
        .iter()
        .map(|d| d as &dyn rusqlite::types::ToSql)
        .collect();

    let rows = stmt
        .query_map(params.as_slice(), |row| {
            Ok((
                row.get::<_, String>(0)?,
                row.get::<_, Vec<u8>>(1)?,
                row.get::<_, String>(2)?,
                row.get::<_, Option<String>>(3)?,
                row.get::<_, bool>(4)?,
                row.get::<_, bool>(5)?,
            ))
        })
        .expect("failed to query cookies");

    let mut result = Vec::new();
    for row in rows {
        let (name, encrypted_value, host_key, path, is_secure, is_httponly) =
            row.expect("failed to read row");

        if let Some(value) = decrypt_v10(&encrypted_value, &key) {
            result.push(Cookie {
                name,
                value,
                domain: host_key,
                path: path.unwrap_or_else(|| "/".to_string()),
                secure: is_secure,
                http_only: is_httponly,
            });
        }
    }

    println!(
        "{}",
        serde_json::to_string(&result).expect("failed to serialize")
    );
}