pass_manager/src/state/get_user.rs

use crate::{errors::HandlerUsed, models::User, 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 {
    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<Bot>, file: &FileMeta) -> crate::Result<Box<[u8]>> {
    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<String>,
    locale: LocaleRef,
) -> crate::Result<Result<(), String>> {
    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.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")
        )?;
    }

    write!(error_text, "\n{}", locale.fix_that_and_send_again)?;

    Ok(Err(error_text))
}

#[inline]
fn user_from_bytes(
    bytes: impl AsRef<[u8]>,
    existing_names: ahash::HashSet<String>,
    locale: LocaleRef,
) -> crate::Result<Result<User, String>> {
    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<Bot>,
    db: &Pool,
    document: Option<&Document>,
    user_id: u64,
    locale: LocaleRef,
) -> Result<User, Cow<'static, str>> {
    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_str());

        let existing_names = Account::get_names(user_id, db)
            .try_collect()
            .map_err(|_| locale.error_getting_account_names.as_str());

        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<Bot>,
    msg: Message,
    db: Pool,
    dialogue: MainDialogue,
    locale: LocaleRef,
    handler: PackagedHandler<User>,
) -> 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,
                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(())
}