use crate::prelude::*; use futures::TryFutureExt; use itertools::Itertools; use std::{fmt::Write, path::Path}; use teloxide::{ net::Download, types::{Document, FileMeta}, }; use tokio::{task::spawn_blocking, try_join}; use trim_in_place::TrimInPlace; #[inline] fn validate_document(document: Option<&Document>) -> Result<&Document, &'static str> { let Some(document) = document else { return Err("You didn't send a file. Try again"); }; if document.file.size > 1024 * 1024 * 200 { return Err("The file is larger than 200 MiB. Try splitting it into multiple files"); } let name = match document.file_name.as_deref() { Some(name) => Path::new(name.trim_end()), None => return Err("Couldn't get the name of the file. Try sending it again"), }; match name.extension() { Some(ext) if ext.eq_ignore_ascii_case("json") => Ok(document), _ => Err("Invalid file name. You need to send a json file. Try again"), } } #[inline] async fn download_file(bot: &Throttle, file: &FileMeta) -> crate::Result> { let path = bot.get_file(file.id.as_str()).await?.path; let mut data = Vec::with_capacity(file.size as usize); bot.download_file_stream(&path) .try_for_each(|bytes| { data.extend_from_slice(&bytes); async { Ok(()) } }) .await?; Ok(data) } #[inline] fn process_accounts( accounts: &mut [DecryptedAccount], existing_names: ahash::HashSet, ) -> 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 = "Your accounts have the following problems:".to_owned(); if !duplicates.is_empty() { write!( error_text, "\n\nDuplicate names:\n{:?}", duplicates.into_iter().format("\n") )?; } if !existing.is_empty() { write!( error_text, "\n\nAccounts with these names already exist in the database:\n{:?}", existing.into_iter().format("\n") )?; } if !invalid.is_empty() { write!( error_text, "\n\nInvalid account fields:\n{:?}", invalid.into_iter().format("\n") )?; } error_text.push_str("\n\nFix these problems and send the file again"); Ok(Err(error_text)) } #[inline] fn user_from_vec( vector: Vec, existing_names: ahash::HashSet, ) -> crate::Result> { let mut user: User = serde_json::from_slice(&vector)?; drop(vector); match process_accounts(&mut user.accounts, existing_names)? { Ok(()) => Ok(Ok(user)), Err(error_text) => Ok(Err(error_text)), } } /// Function to handle `GetUser` state. It doesn't actually validate anything pub async fn get_user( bot: Throttle, msg: Message, db: Pool, dialogue: MainDialogue, 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, "Successfully cancelled", deletion_markup(), None) .await?; return Ok(()); } let file = match validate_document(msg.document()) { Ok(document) => &document.file, Err(text) => { handler .previous .alter_message(&bot, text, None, None) .await?; return Ok(()); } }; let existing_names = Account::get_names(user_id, &db) .try_collect() .map_err(Into::into); let (data, existing_names) = try_join!(download_file(&bot, file), existing_names)?; let user = match spawn_blocking(move || user_from_vec(data, existing_names)).await? { Ok(Ok(user)) => user, Ok(Err(error_text)) => { handler .previous .alter_message(&bot, error_text, None, None) .await?; return Ok(()); } Err(_) => { handler .previous .alter_message(&bot, "Error parsing the json file. Try again", 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, user).await { let _ = dialogue.exit().await; return Err(err); } Ok(()) }