use crate::prelude::*; use futures::TryFutureExt; use itertools::Itertools; use std::{borrow::Cow, fmt::Write, path::Path}; use teloxide::{ net::Download, types::{Document, FileMeta}, }; use tokio::{task::spawn_blocking, try_join}; use trim_in_place::TrimInPlace; #[derive(Clone, Copy)] enum InvalidDocument { NoFileSend, FileTooLarge, CouldntGetFileName, InvalidFileName, } impl InvalidDocument { const fn into_str(self, locale: LocaleRef) -> &'static str { match self { Self::NoFileSend => &locale.no_file_send, Self::FileTooLarge => &locale.file_too_large, Self::CouldntGetFileName => &locale.couldnt_get_file_name, Self::InvalidFileName => &locale.invalid_file_name, } } } #[inline] fn validate_document(document: Option<&Document>) -> Result<&Document, InvalidDocument> { let Some(document) = document else { return Err(InvalidDocument::NoFileSend); }; if document.file.size > 1024 * 1024 * 200 { return Err(InvalidDocument::FileTooLarge); } let name = match document.file_name.as_deref() { Some(name) => Path::new(name.trim_end()), None => return Err(InvalidDocument::CouldntGetFileName), }; match name.extension() { Some(ext) if ext.eq_ignore_ascii_case("json") => Ok(document), _ => Err(InvalidDocument::InvalidFileName), } } #[inline] async fn download_file(bot: &Throttle, file: &FileMeta) -> crate::Result> { let path = bot.get_file(&file.id).await?.path; let mut data = Vec::with_capacity(file.size as usize); let mut stream = bot.download_file_stream(&path); while let Some(bytes) = stream.try_next().await? { data.extend_from_slice(&bytes); } Ok(data.into_boxed_slice()) } #[inline] fn process_accounts( accounts: &mut [DecryptedAccount], existing_names: ahash::HashSet, locale: LocaleRef, ) -> crate::Result> { for account in accounts.iter_mut() { account.name.trim_in_place(); account.login.trim_in_place(); account.password.trim_in_place(); } accounts.sort_unstable_by(|a, b| a.name.cmp(&b.name)); let mut duplicates = Vec::new(); let mut existing = Vec::new(); let mut invalid = Vec::new(); accounts .iter() .dedup_by_with_count(|a, b| a.name == b.name) .for_each(|(count, account)| { let duplicate = count != 1; let exists = existing_names.contains(&account.name); if duplicate { duplicates.push(account.name.as_str()); } if exists { existing.push(account.name.as_str()); } // If it already exists or if it is a duplicate there's no need to check the account's validity if !duplicate && !exists && !account.validate() { invalid.push(account.name.as_str()); } }); drop(existing_names); if duplicates.is_empty() && invalid.is_empty() && existing.is_empty() { return Ok(Ok(())); } let mut error_text = locale.following_accounts_have_problems.as_ref().to_owned(); if !duplicates.is_empty() { write!( error_text, "\n\n{}:\n{:?}", locale.duplicate_names, duplicates.into_iter().format("\n") )?; } if !existing.is_empty() { write!( error_text, "\n\n{}:\n{:?}", locale.accounts_already_in_db, existing.into_iter().format("\n") )?; } if !invalid.is_empty() { write!( error_text, "\n\n{}:\n{:?}", locale.invalid_fields, invalid.into_iter().format("\n") )?; } error_text.push_str(&locale.fix_that_and_send_again); Ok(Err(error_text)) } #[inline] fn user_from_bytes( bytes: impl AsRef<[u8]>, existing_names: ahash::HashSet, locale: LocaleRef, ) -> crate::Result> { let mut user: User = serde_json::from_slice(bytes.as_ref())?; drop(bytes); match process_accounts(&mut user.accounts, existing_names, locale)? { Ok(()) => Ok(Ok(user)), Err(error_text) => Ok(Err(error_text)), } } #[inline] async fn user_from_document( bot: &Throttle, db: &Pool, document: Option<&Document>, user_id: u64, locale: LocaleRef, ) -> Result> { let (data, existing_names) = { let file = &validate_document(document) .map_err(|err| err.into_str(locale))? .file; let data = download_file(bot, file).map_err(|_| locale.error_downloading_file.as_ref()); let existing_names = Account::get_names(user_id, db) .try_collect() .map_err(|_| locale.error_getting_account_names.as_ref()); try_join!(data, existing_names)? }; match spawn_blocking(|| user_from_bytes(data, existing_names, locale)).await { Ok(Ok(user)) => user.map_err(Cow::Owned), _ => Err(Cow::Borrowed(&locale.error_parsing_json_file)), } } /// Function to handle `GetUser` state. It doesn't actually validate anything #[inline] pub async fn get_user( bot: Throttle, msg: Message, db: Pool, dialogue: MainDialogue, locale: LocaleRef, handler: PackagedHandler, ) -> crate::Result<()> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let mut handler = handler.lock().await; if handler.func.is_none() { let _ = dialogue.exit().await; return Err(HandlerUsed.into()); } if msg.text().map(str::trim) == Some("/cancel") { dialogue.exit().await?; handler .previous .alter_message( &bot, locale.successfully_canceled.as_ref(), deletion_markup(locale), None, ) .await?; return Ok(()); } let user = match user_from_document(&bot, &db, msg.document(), user_id, locale).await { Ok(user) => user, Err(error_text) => { handler .previous .alter_message(&bot, error_text, None, None) .await?; return Ok(()); } }; let previous = handler.previous; let func = handler.func.take().unwrap(); drop(handler); if let Err(err) = func(bot, msg, db, dialogue.clone(), previous, locale, user).await { let _ = dialogue.exit().await; return Err(err); } Ok(()) }