hermes-tdata 0.1.1

//! MTP authorization data structures and parsing
//!
//! Handles reading and decrypting MTP data files containing auth keys.

use std::path::Path;

use crate::crypto::{AuthKey, decrypt_local};
use crate::qdatastream::QDataStream;
use crate::{Error, Result};

use super::file_io::read_file;

/// MTP authorization data
#[derive(Debug)]
pub struct MtpData {
    pub dc_id: i32,
    pub user_id: i64,
    pub auth_key: [u8; 256],
}

/// Special tag for wide (64-bit) user IDs
const K_WIDE_IDS_TAG: i64 = !0i64; // All bits set = -1

/// Read MTP data file (contains the actual auth key)
///
/// The MTP data is stored in a file named by ToFilePart(ComputeDataNameKey(keyFile))
/// where keyFile is like "data" or "data#1" for multi-account
pub fn read_mtp_data(
    base_path: &Path,
    index: i32,
    local_key: &AuthKey,
    key_file: &str,
) -> Result<MtpData> {
    // Compute data name key = MD5(keyFile)
    let data_name = compose_data_string(key_file, index);
    let data_name_key = compute_data_name_key(&data_name);
    let file_name = to_file_part(data_name_key);

    tracing::debug!("Looking for MTP data in file: {}", file_name);

    // Read the encrypted file
    let file = read_file(&file_name, base_path)?;

    // The file contains a QByteArray which is the encrypted data
    let mut stream = QDataStream::new(&file.data);
    let encrypted = stream.read_qbytearray()?;

    // Decrypt
    let decrypted = decrypt_local(&encrypted, local_key)?;

    // Parse the decrypted MTP data
    parse_mtp_authorization(&decrypted)
}

/// Compose data string: "data" for index 0, "data#2" for index 1, etc.
fn compose_data_string(key_file: &str, index: i32) -> String {
    let base = key_file.replace('#', "");
    if index > 0 { format!("{}#{}", base, index + 1) } else { base }
}

/// Compute data name key from key file name using MD5
fn compute_data_name_key(data_name: &str) -> u64 {
    use md5::{Digest, Md5};

    let mut hasher = Md5::new();
    hasher.update(data_name.as_bytes());
    let result: [u8; 16] = hasher.finalize().into();

    // Take lower 64 bits (little endian)
    u64::from_le_bytes([
        result[0], result[1], result[2], result[3], result[4], result[5], result[6], result[7],
    ])
}

/// Convert a `FileKey` (u64) to a 16-character hex file name
fn to_file_part(val: u64) -> String {
    let mut result = String::with_capacity(16);
    let mut v = val;

    for _ in 0..16 {
        let nibble = (v & 0x0F) as u8;
        let c =
            if nibble < 0x0A { (b'0' + nibble) as char } else { (b'A' + (nibble - 0x0A)) as char };
        result.push(c);
        v >>= 4;
    }

    result
}

/// Parse MTP authorization data from decrypted bytes
///
/// Format:
/// - int32: `block_id` (must be 0x4B = dbi.MtpAuthorization)
/// - `QByteArray`: serialized authorization data
///
/// Serialized format:
/// - int32: userId (or kWideIdsTag for new format)
/// - int32: mainDcId (or if kWideIdsTag: int64 userId, int32 mainDcId)
/// - int32: keysCount
/// - for each key:
///   - int32: dcId
///   - bytes[256]: authKey
/// - int32: keysToDestroyCount
/// - ...
fn parse_mtp_authorization(data: &[u8]) -> Result<MtpData> {
    let mut stream = QDataStream::new(data);

    // Read block ID
    let block_id = stream.read_i32()?;

    // 0x4B = 75 = dbi.MtpAuthorization
    if block_id != 0x4B {
        return Err(Error::invalid_format(format!(
            "expected MtpAuthorization block (0x4B), got 0x{block_id:02X}"
        )));
    }

    // Read the serialized QByteArray
    let serialized = stream.read_qbytearray()?;
    let mut auth_stream = QDataStream::new(&serialized);

    // Read user ID and DC ID
    let first_int = auth_stream.read_i32()?;
    let second_int = auth_stream.read_i32()?;

    // Check for kWideIdsTag (new format with 64-bit user ID)
    let combined = (i64::from(first_int) << 32) | i64::from(second_int as u32);

    let (user_id, main_dc_id) = if combined == K_WIDE_IDS_TAG {
        // New format: next is int64 userId, then int32 mainDcId
        let uid = auth_stream.read_i64()?;
        let dc = auth_stream.read_i32()?;
        (uid, dc)
    } else {
        // Old format: first_int is userId, second_int is mainDcId
        (i64::from(first_int), second_int)
    };

    tracing::debug!("MTP auth: user_id={}, main_dc_id={}", user_id, main_dc_id);

    // Read keys count
    let keys_count = auth_stream.read_i32()?;

    if !(0..=10).contains(&keys_count) {
        return Err(Error::invalid_format(format!("invalid keys count: {keys_count}")));
    }

    // Read auth keys
    let mut auth_key: Option<[u8; 256]> = None;

    for _ in 0..keys_count {
        let dc_id = auth_stream.read_i32()?;
        let key_bytes = auth_stream.read_raw(256)?;

        tracing::debug!("Found key for DC {}", dc_id);

        if dc_id == main_dc_id {
            let mut key = [0u8; 256];
            key.copy_from_slice(&key_bytes);
            auth_key = Some(key);
        }
    }

    let auth_key = auth_key.ok_or_else(|| {
        Error::auth_key_failed(format!("no auth key found for main DC {main_dc_id}"))
    })?;

    Ok(MtpData { dc_id: main_dc_id, user_id, auth_key })
}