use crate::prelude::*; use itertools::Itertools; use rustc_hash::FxHashSet; 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 document = match document { Some(document) => document, None => 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(bytes); async { Ok(()) } }) .await?; Ok(data) } #[inline] fn process_accounts( accounts: &mut [DecryptedAccount], existing_names: FxHashSet, ) -> 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: FxHashSet, ) -> 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: DatabaseConnection, dialogue: MainDialogue, next: PackagedHandler, ) -> crate::Result<()> { let user_id = msg.from().ok_or(NoUserInfo)?.id.0; let mut handler = next.lock().await; if handler.func.is_none() { let _ = dialogue.exit().await; return Err(HandlerUsed.into()); } handler.previous.delete(&bot).await; if let Some("/cancel") = msg.text().map(str::trim) { dialogue.exit().await?; bot.send_message(msg.chat.id, "Successfully cancelled") .reply_markup(deletion_markup()) .await?; return Ok(()); } let file = match validate_document(msg.document()) { Ok(document) => &document.file, Err(text) => { let msg = bot.send_message(msg.chat.id, text).await?; handler.previous = MessageIds::from(&msg); return Ok(()); } }; let existing_names = async { Account::get_names(user_id, &db) .await? .try_collect::>() .await .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)) => { let msg = bot.send_message(msg.chat.id, error_text).await?; handler.previous = MessageIds::from(&msg); return Ok(()); } Err(_) => { let msg = bot .send_message(msg.chat.id, "Error parsing the json file. Try again") .await?; handler.previous = MessageIds::from(&msg); return Ok(()); } }; let func = handler.func.take().unwrap(); drop(handler); if let Err(err) = func(bot, msg, db, dialogue.clone(), user).await { let _ = dialogue.exit().await; return Err(err); } Ok(()) }