sdgb-utils-rs/sdgb-cli/src/main.rs

use std::sync::atomic::{AtomicBool, Ordering};

use nyquest_preset::nyquest::ClientBuilder;
use palc::Parser;
use spdlog::{Level, LevelFilter::MoreSevereEqual};

use sdgb_api::{
    all_net::QRCode,
    auth_lite::{SDGB, SDHJ, delivery_raw},
    title::{
        MaiVersionExt, Sdgb1_40, Sdgb1_50,
        methods::APIMethod,
        model::{
            GetUserDataApi, GetUserDataApiResp, GetUserPreviewApi, GetUserPreviewApiResp, Ping,
            PingResp, UserLogoutApi, UserLogoutApiResp,
        },
    },
};
use spdlog::{error, info, warn};

use crate::{commands::Cli, utils::login_action};

#[cfg(feature = "fetchall")]
mod cache;
mod commands;
mod utils;

static EARLY_QUIT: AtomicBool = AtomicBool::new(false);

#[cfg_attr(feature = "compio", compio::main)]
#[cfg_attr(feature = "tokio", tokio::main)]
async fn main() -> Result<(), Box<dyn snafu::Error>> {
    nyquest_preset::register();

    if cfg!(debug_assertions) {
        spdlog::default_logger().set_level_filter(MoreSevereEqual(Level::Debug));
    } else {
        spdlog::default_logger().set_level_filter(MoreSevereEqual(Level::Info));
    }

    ctrlc::set_handler(|| {
        if EARLY_QUIT.load(Ordering::Relaxed) {
            error!("force-quit triggered!");
            std::process::exit(1);
        } else {
            warn!("received early-quit request! will abort soon");
            EARLY_QUIT.store(true, Ordering::Relaxed);
        }
    })?;

    let Cli { command } = <Cli as Parser>::parse();

    let client = ClientBuilder::default().build_async().await?;

    // TODO: refactor via enum_dispatch
    match command {
        commands::Commands::Logout { user_id } => {
            let logout: UserLogoutApiResp = Sdgb1_50::request(
                &client,
                APIMethod::UserLogoutApi,
                user_id,
                UserLogoutApi {
                    user_id,
                    ..Default::default()
                },
            )
            .await?;
            println!("{logout:?}");
        }
        commands::Commands::Preview { user_id } => {
            let preview: GetUserPreviewApiResp = Sdgb1_50::request(
                &client,
                APIMethod::GetUserPreviewApi,
                user_id,
                GetUserPreviewApi { user_id },
            )
            .await?;

            println!("{preview}");
        }
        commands::Commands::Ping => {
            let decoded: PingResp = Sdgb1_40::request(
                &client,
                APIMethod::Ping,
                "",
                Ping {}, // note: must not be `Ping`, or serde_json serializes to nothing
            )
            .await?;
            info!("sdgb 1.40 resp: {decoded}");
            let decoded: PingResp =
                Sdgb1_50::request(&client, APIMethod::Ping, "", Ping {}).await?;
            info!("sdgb 1.50 resp: {decoded}");
        }
        commands::Commands::QRLogin { ref qrcode_content } => {
            let qrcode = QRCode { qrcode_content };
            match qrcode.login(&client).await {
                Ok(user_id) => info!("login succeed: {user_id}"),
                Err(e) => error!("login failed: {e}"),
            }
        }

        commands::Commands::AuthLite { title_ver, variant } => {
            let resp = match variant {
                commands::AuthLiteVariant::SDGB => delivery_raw::<SDGB>(&client, title_ver).await?,
                commands::AuthLiteVariant::SDHJ => delivery_raw::<SDHJ>(&client, title_ver).await?,
            };
            println!("{}", String::from_utf8_lossy(&resp));
        }

        #[cfg(feature = "fetchall")]
        commands::Commands::ListAllUser { concurrency } => {
            use futures_util::StreamExt;
            use sdgb_api::bincode::borrow_decode_from_slice;
            use std::io::{self, BufRead};

            let mut user_ids = Vec::new();
            {
                let mut stdin = io::stdin().lock();
                let mut buf = String::new();

                while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) {
                    if buf.is_empty() {
                        continue;
                    }

                    let user_id: u32 = buf.trim().parse()?;
                    buf.clear();
                    user_ids.push(user_id);
                }
            }

            let _ = cache::init_db();
            let read = cache::read_txn()?;
            let write = cache::write_txn()?;
            let config = sdgb_api::bincode::config::Configuration::<
                sdgb_api::bincode::config::LittleEndian,
            >::default()
            .with_no_limit();

            info!("number of user_id: {}", user_ids.len());

            let collect = futures_util::stream::iter(user_ids)
                .map(async |user_id| {
                    {
                        let cache_table = cache::open_table_read(&read)?;
                        let data = cache_table.get(user_id)?;
                        if let Some(data) = data {
                            let decoded: (GetUserPreviewApiResp, _) =
                                borrow_decode_from_slice(&data.value(), config)?;

                            return Ok(decoded.0);
                        }
                    }

                    if EARLY_QUIT.load(Ordering::Relaxed) {
                        return Err("early skip due to ctrl-c")?;
                    }

                    let resp = Sdgb1_50::request::<_, GetUserPreviewApiResp>(
                        &client,
                        APIMethod::GetUserPreviewApi,
                        user_id,
                        GetUserPreviewApi { user_id },
                    )
                    .await;

                    match &resp {
                        Ok(resp) => {
                            use sdgb_api::bincode::encode_to_vec;

                            info!("found: {user_id}");

                            if let Ok(mut table) = cache::open_table(&write)
                                && let Ok(encoded) = encode_to_vec(resp, config)
                            {
                                _ = table.insert(resp.user_id, encoded);
                            }
                        }
                        Err(sdgb_api::ApiError::JSON { .. }) => {}
                        Err(e) => {
                            error!("preview failed: {e}");
                        }
                    }

                    Result::<_, Box<dyn snafu::Error>>::Ok(resp?)
                })
                .buffer_unordered(concurrency) // slower to avoid being banned
                .filter_map(async |r| r.ok())
                .collect::<Vec<_>>()
                .await;
            drop(collect);

            let _ = write.commit();
        }
        #[cfg(feature = "fetchall")]
        commands::Commands::ListAllUserDump { .. } => {
            use std::{fs::OpenOptions, io::BufWriter};

            use redb::ReadableTable;
            use sdgb_api::bincode::{self, borrow_decode_from_slice};

            use crate::cache::{open_table_read, read_txn};

            let txn = read_txn()?;
            let table = open_table_read(&txn)?;

            let config = bincode::config::Configuration::<bincode::config::LittleEndian>::default()
                .with_no_limit();

            let user_ids = table
                .iter()?
                .flatten()
                .map(|d| borrow_decode_from_slice(&d.1.value(), config))
                .flatten()
                .map(|(value, _)| value)
                .collect::<Vec<GetUserPreviewApiResp>>();

            let file = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .open("players.json")?;

            #[rustversion::since(1.89)]
            {
                file.lock()?;
            }

            let writer = BufWriter::new(file);
            serde_json::to_writer(writer, &user_ids)?;
            info!("dumped {} user id", user_ids.len());
        }

        commands::Commands::Userdata { user_id } => {
            let action = async |_| match Sdgb1_50::request::<_, GetUserDataApiResp>(
                &client,
                APIMethod::GetUserDataApi,
                user_id,
                GetUserDataApi { user_id },
            )
            .await
            {
                Ok(udata) => {
                    info!("{udata:#?}");
                }
                Err(e) => {
                    error!("failed to get userdata: {e}");
                }
            };
            login_action(&client, user_id, action).await?;
        }
    }

    Ok(())
}

#[cfg(all(feature = "compio", feature = "tokio"))]
compile_error!("you must not enable both `compio` and `tokio`");