use std::path::Path;
use md5::{Digest, Md5};
use crate::crypto::{self, AuthKey};
use crate::qt::{QReader, QWriter};
use crate::tdata::map::MapData;
use crate::tdata::mtp_config::{DcOptions, Environment, MtpConfig};
use crate::tdata::{FillRandom, encrypt, hex_bytes, write_tdf};
use crate::{Error, MagicContext, Result, tdf};
const MTP_MAIN_FIELDS_FAILED: &str = "Could not read main fields from mtp authorization.";
const MTP_KEY_COUNT_FAILED: &str = "Could not read keys count from mtp authorization.";
pub const DBI_MTP_AUTHORIZATION: u32 = 0x4B;
pub const DBI_KEY: u32 = 0x00;
pub const DBI_USER: u32 = 0x01;
pub const DBI_DC_OPTION_OLD_OLD: u32 = 0x02;
pub const DBI_CHAT_SIZE_MAX_OLD: u32 = 0x03;
pub const DBI_DC_OPTION_OLD: u32 = 0x27;
pub const DBI_MEGAGROUP_SIZE_MAX_OLD: u32 = 0x32;
pub const DBI_SAVED_GIFS_LIMIT_OLD: u32 = 0x35;
pub const DBI_STICKERS_RECENT_LIMIT_OLD: u32 = 0x43;
pub const DBI_DC_OPTIONS_OLD: u32 = 0x4A;
pub const DBI_STICKERS_FAVED_LIMIT_OLD: u32 = 0x50;
pub const DBI_APPLICATION_SETTINGS: u32 = 0x5E;
pub const MTP_CONFIG_PRODUCTION: [u8; 748] = hex_bytes(concat!(
"0000000100000000000001fefffffffe0000000b0000000100000010000001bb",
"0000000e3134392e3135342e3137352e35300000000000000001000000110000",
"01bb00000027323030313a306232383a663233643a663030313a303030303a30",
"3030303a303030303a30303061000000000000000200000010000001bb000000",
"0e3134392e3135342e3136372e3531000000000000000200000010000001bb00",
"00000d39352e3136312e37362e313030000000000000000200000011000001bb",
"00000027323030313a303637633a303465383a663030323a303030303a303030",
"303a303030303a30303061000000000000000300000010000001bb0000000f31",
"34392e3135342e3137352e313030000000000000000300000011000001bb0000",
"0027323030313a306232383a663233643a663030333a303030303a303030303a",
"303030303a30303061000000000000000400000010000001bb0000000e313439",
"2e3135342e3136372e3931000000000000000400000011000001bb0000002732",
"3030313a303637633a303465383a663030343a303030303a303030303a303030",
"303a30303061000000000000000500000010000001bb0000000d3134392e3135",
"342e3137312e35000000000000000500000011000001bb00000027323030313a",
"306232383a663233663a663030353a303030303a303030303a303030303a3030",
"30610000000000000000000000c800002710000000640001d4c0000013880000",
"7530000003e8000493e000007530000005dc000000c80002a3000002a3000002",
"a300000000000000001e0000000500000005000000640000001a006800740074",
"00700073003a002f002f0074002e006d0065002f00093a8000004e2000015f90",
"0000753000002710000000040000001a0061007000760033002e007300740065",
"006c002e0063006f006d00000001000000000000040000000004d83ddc4d0000",
"0000000000000024ea0000000006006700690066000000140066006f00750072",
"007300710075006100720065",
));
pub const WIDE_IDS_TAG: i64 = -1;
pub const CONFIG_FILE: &str = "config";
#[must_use]
pub fn data_name(key_file: &str, index: i32) -> String {
let name = key_file.replace('#', "");
if index > 0 {
format!("{name}#{}", i64::from(index) + 1)
} else {
name
}
}
#[must_use]
pub fn data_name_key(data_name: &str) -> u64 {
let digest = Md5::digest(data_name.as_bytes());
let mut low = [0; 8];
low.copy_from_slice(&digest[..8]);
u64::from_le_bytes(low)
}
#[must_use]
pub fn file_part(key: u64) -> String {
(0..16)
.map(|i| {
let nibble = u32::try_from((key >> (4 * i)) & 0xF).unwrap_or_default();
char::from_digit(nibble, 16)
.unwrap_or('0')
.to_ascii_uppercase()
})
.collect()
}
#[must_use]
pub fn folder_name(data_name: &str) -> String {
file_part(data_name_key(data_name))
}
fn read_encrypted(dir: &Path, name: &str, local_key: &AuthKey) -> Result<Vec<u8>> {
let file = tdf::read(dir, name)?;
let encrypted = QReader::new(&file.payload)
.read_bytearray()
.ok()
.flatten()
.unwrap_or_default();
crypto::decrypt_local(encrypted, local_key)
}
fn write_encrypted(
dir: &Path,
name: &str,
data: &[u8],
local_key: &AuthKey,
rng: &mut FillRandom<'_>,
) -> Result<()> {
let mut payload = QWriter::new();
payload.write_bytearray(Some(&encrypt(data, local_key, rng)?))?;
write_tdf(dir, name, payload.into_bytes())
}
struct RawReader<'a> {
rest: &'a [u8],
}
impl<'a> RawReader<'a> {
fn with<T>(&mut self, read: impl FnOnce(&mut QReader<'a>) -> Result<T>) -> Result<T> {
let mut reader = QReader::new(self.rest);
let value = read(&mut reader)?;
let used = self.rest.len() - reader.remaining();
self.rest = &self.rest[used..];
Ok(value)
}
fn i32(&mut self) -> Result<i32> {
self.with(QReader::read_i32)
}
fn key(&mut self) -> Result<AuthKey> {
let bytes = self
.rest
.get(..AuthKey::SIZE)
.ok_or(Error::QDataStreamFailed { context: None })?;
self.rest = &self.rest[AuthKey::SIZE..];
AuthKey::from_slice(bytes)
}
fn main_fields(&mut self) -> Result<(u64, i32)> {
let legacy_user = self.i32()?;
let legacy_dc = self.i32()?;
if (i64::from(legacy_user) << 32) | i64::from(legacy_dc) == WIDE_IDS_TAG {
Ok((self.with(QReader::read_u64)?, self.i32()?))
} else {
Ok((i64::from(legacy_user).cast_unsigned(), legacy_dc))
}
}
fn keys(&mut self) -> Result<Vec<(i32, AuthKey)>> {
let count = self.i32().map_err(|_| Error::QDataStreamFailed {
context: Some(MTP_KEY_COUNT_FAILED),
})?;
let mut keys = Vec::new();
for _ in 0..count {
let dc = self.i32()?;
keys.push((dc, self.key()?));
}
Ok(keys)
}
}
#[derive(Default)]
struct MtpData<'a> {
authorization: &'a [u8],
legacy_user_id: i32,
legacy_main_dc: u32,
legacy_keys: Vec<(i32, AuthKey)>,
}
impl<'a> MtpData<'a> {
fn parse(data: &'a [u8]) -> Result<Self> {
let mut reader = RawReader { rest: data };
let mut out = Self::default();
while !reader.rest.is_empty() {
let block_id =
reader
.with(QReader::read_u32)
.map_err(|_| Error::QDataStreamFailed {
context: Some("Could not read a block id from MtpData"),
})?;
match block_id {
DBI_MTP_AUTHORIZATION => {
out.authorization = reader.with(QReader::read_bytearray)?.unwrap_or_default();
}
DBI_KEY => {
let dc = reader.i32()?;
out.legacy_keys.push((dc, reader.key()?));
}
DBI_USER => {
out.legacy_user_id = reader.i32()?;
out.legacy_main_dc = reader.with(QReader::read_u32)?;
}
DBI_DC_OPTION_OLD_OLD => {
reader.with(|r| {
r.read_u32()?;
r.read_qstring()?;
r.read_qstring()?;
r.read_u32()
})?;
}
DBI_DC_OPTION_OLD => {
reader.with(|r| {
r.read_u32()?;
r.read_i32()?;
r.read_qstring()?;
r.read_u32()
})?;
}
DBI_DC_OPTIONS_OLD => {
let blob = reader.with(QReader::read_bytearray)?.unwrap_or_default();
DcOptions::parse(blob)?;
}
DBI_APPLICATION_SETTINGS => {
reader.with(QReader::read_bytearray)?;
}
DBI_CHAT_SIZE_MAX_OLD
| DBI_SAVED_GIFS_LIMIT_OLD
| DBI_STICKERS_RECENT_LIMIT_OLD
| DBI_STICKERS_FAVED_LIMIT_OLD
| DBI_MEGAGROUP_SIZE_MAX_OLD => {
reader.i32()?;
}
_ => {
return Err(Error::InvalidMagic {
context: MagicContext::MtpDataBlock { block_id },
});
}
}
}
Ok(out)
}
}
fn write_keys(out: &mut Vec<u8>, keys: &[(i32, AuthKey)]) -> Result<()> {
let count =
i32::try_from(keys.len()).map_err(|_| Error::QDataStreamFailed { context: None })?;
out.extend_from_slice(&count.to_be_bytes());
for (dc, key) in keys {
out.extend_from_slice(&dc.to_be_bytes());
out.extend_from_slice(key.as_bytes());
}
Ok(())
}
#[derive(Debug, Clone)]
pub struct Account {
pub index: i32,
pub user_id: u64,
pub main_dc: i32,
pub keys: Vec<(i32, AuthKey)>,
pub keys_to_destroy: Vec<(i32, AuthKey)>,
pub map: MapData,
pub config: Option<Vec<u8>>,
}
impl Account {
pub fn new(index: i32, user_id: u64, main_dc: i32, keys: Vec<(i32, AuthKey)>) -> Result<Self> {
let account = Self {
index,
user_id,
main_dc,
keys,
keys_to_destroy: Vec::new(),
map: MapData::new_account(),
config: None,
};
if account.auth_key().is_none() {
return Err(Error::AuthKeyNotFound);
}
Ok(account)
}
pub fn parse_mtp_authorization(index: i32, serialized: &[u8]) -> Result<Self> {
let mut reader = RawReader { rest: serialized };
let (user_id, main_dc) = reader.main_fields().map_err(|_| Error::QDataStreamFailed {
context: Some(MTP_MAIN_FIELDS_FAILED),
})?;
let keys = reader.keys()?;
let keys_to_destroy = reader.keys()?;
let account = Self {
index,
user_id,
main_dc,
keys,
keys_to_destroy,
map: MapData::default(),
config: None,
};
if account.auth_key().is_none() {
return Err(Error::AuthKeyNotFound);
}
Ok(account)
}
pub fn encode_mtp_authorization(&self) -> Result<Vec<u8>> {
let mut out = Vec::new();
out.extend_from_slice(&WIDE_IDS_TAG.to_be_bytes());
out.extend_from_slice(&self.user_id.to_be_bytes());
out.extend_from_slice(&self.main_dc.to_be_bytes());
write_keys(&mut out, &self.keys)?;
write_keys(&mut out, &self.keys_to_destroy)?;
Ok(out)
}
pub fn encode_mtp_data(&self) -> Result<Vec<u8>> {
let mut writer = QWriter::new();
writer.write_u32(DBI_MTP_AUTHORIZATION);
writer.write_bytearray(Some(&self.encode_mtp_authorization()?))?;
Ok(writer.into_bytes())
}
pub(crate) fn write(
&self,
dir: &Path,
key_file: &str,
local_key: &AuthKey,
performance_mode: bool,
rng: &mut FillRandom<'_>,
) -> Result<()> {
let part = folder_name(&data_name(key_file, self.index));
let folder = dir.join(&part);
self.map.write(&folder, local_key, rng)?;
if !performance_mode {
write_encrypted(&folder, CONFIG_FILE, &self.config_data()?, local_key, rng)?;
}
write_encrypted(dir, &part, &self.encode_mtp_data()?, local_key, rng)
}
pub(crate) fn write_plaintexts(&self, performance_mode: bool) -> Result<Vec<Vec<u8>>> {
let mut out = vec![self.map.plaintext()?];
if !performance_mode {
out.push(self.config_data()?);
}
out.push(self.encode_mtp_data()?);
Ok(out)
}
fn config_data(&self) -> Result<Vec<u8>> {
let config = self.config.as_deref().unwrap_or(&MTP_CONFIG_PRODUCTION);
let mut data = QWriter::new();
data.write_bytearray(Some(config))?;
Ok(data.into_bytes())
}
#[must_use]
pub fn auth_key(&self) -> Option<&AuthKey> {
self.keys
.iter()
.find(|(dc, _)| *dc == self.main_dc)
.map(|(_, key)| key)
}
pub fn parse_mtp_data(index: i32, data: &[u8]) -> Result<Self> {
let mtp = MtpData::parse(data)?;
if !mtp.authorization.is_empty() {
return Self::parse_mtp_authorization(index, mtp.authorization);
}
if mtp.legacy_keys.is_empty() && mtp.legacy_main_dc == 0 {
return Err(Error::AuthKeyNotFound);
}
let account = Self {
index,
user_id: i64::from(mtp.legacy_user_id).cast_unsigned(),
main_dc: mtp.legacy_main_dc.cast_signed(),
keys: mtp.legacy_keys,
keys_to_destroy: Vec::new(),
map: MapData::default(),
config: None,
};
if account.auth_key().is_none() {
return Err(Error::AuthKeyNotFound);
}
Ok(account)
}
fn read_mtp_data(dir: &Path, data_name: &str, local_key: &AuthKey) -> Result<Vec<u8>> {
read_encrypted(dir, &folder_name(data_name), local_key)
}
pub fn read_mtp_authorization(
dir: impl AsRef<Path>,
data_name: &str,
local_key: &AuthKey,
) -> Result<Vec<u8>> {
let data = Self::read_mtp_data(dir.as_ref(), data_name, local_key)?;
let mtp = MtpData::parse(&data)?;
if mtp.authorization.is_empty() {
return Err(Error::AuthKeyNotFound);
}
Ok(mtp.authorization.to_vec())
}
#[must_use]
pub fn mtp_config(&self) -> MtpConfig {
self.config
.as_deref()
.and_then(|serialized| MtpConfig::parse(serialized).ok())
.unwrap_or_else(|| MtpConfig::new(Environment::Production))
}
pub fn read(
dir: impl AsRef<Path>,
key_file: &str,
index: i32,
local_key: &AuthKey,
) -> Result<Self> {
let dir = dir.as_ref();
let name = data_name(key_file, index);
let folder = dir.join(folder_name(&name));
let map = MapData::read(&folder, local_key)?;
let data = Self::read_mtp_data(dir, &name, local_key)?;
let mut account = Self::parse_mtp_data(index, &data)?;
account.map = map;
account.config = read_encrypted(&folder, CONFIG_FILE, local_key)
.ok()
.and_then(|data| {
QReader::new(&data)
.read_bytearray()
.ok()
.flatten()
.filter(|blob| !blob.is_empty())
.map(<[u8]>::to_vec)
});
Ok(account)
}
}