use std::{ sync::{ Arc, atomic::{AtomicBool, Ordering}, }, time::SystemTime, }; use nyquest_preset::nyquest::ClientBuilder; use palc::Parser; use spdlog::{Level, LevelFilter::MoreSevereEqual, sink::StdStreamSink, terminal_style::StyleMode}; use sdgb_api::{ all_net::QRCode, auth_lite::{SDGB, SDHJ, delivery_raw}, helper::MUSIC_DB, title::{ MaiVersionExt, Sdgb1_50, methods::APIMethod, model::{ DataVersion, DxCalculatedEntries, DxRatingNet, GetUserDataApi, GetUserDataApiResp, GetUserMusicApi, GetUserMusicApiResp, GetUserPreviewApi, GetUserPreviewApiResp, GetUserRatingApi, GetUserRatingApiResp, Ping, PingResp, UserLogoutApi, UserLogoutApiResp, }, }, }; use spdlog::{error, info, warn}; use crate::{ commands::{Cli, RatingFormat}, utils::{human_readable_display, json_display, 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> { nyquest_preset::register(); let logger = spdlog::default_logger().fork_with(|log| { log.set_level_filter(MoreSevereEqual(if cfg!(debug_assertions) { Level::Debug } else { Level::Info })); let sink = StdStreamSink::builder() .stderr() .style_mode(StyleMode::Auto) .build()?; *log.sinks_mut() = vec![Arc::new(sink)]; Ok(()) })?; spdlog::swap_default_logger(logger); 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 _ = &*MUSIC_DB; let Cli { command, machine_readable, } = ::parse(); let human_readable = !machine_readable; let client = ClientBuilder::default().build_async().await?; // TODO: refactor via enum_dispatch match command { commands::Commands::MusicDetail { user_id } => { let mut music_detail = Vec::new(); let mut index = None; loop { let GetUserMusicApiResp { next_index, mut user_music_list, .. } = Sdgb1_50::request::<_, GetUserMusicApiResp>( &client, APIMethod::GetUserMusicApi, user_id, GetUserMusicApi { user_id, next_index: index.unwrap_or_default(), max_count: 50, }, ) .await?; for list in &mut user_music_list { music_detail.append(&mut list.user_music_detail_list); } if next_index == 0 || user_music_list.is_empty() { break; } index = Some(next_index); } // TODO: `Display` support for MusicDetail json_display(music_detail)?; } commands::Commands::Rating { user_id, format } => { let rating: GetUserRatingApiResp = Sdgb1_50::request( &client, APIMethod::GetUserRatingApi, user_id, GetUserRatingApi { user_id }, ) .await?; match (human_readable, format) { (true, _) => println!("{rating}"), (false, RatingFormat::Origin) => json_display(rating)?, (false, RatingFormat::DxRatingNet) => { let mut data = DxCalculatedEntries::from_user_rating_lossy(&rating.user_rating); let mut records = data.b35; records.append(&mut data.b15); json_display(records)?; } (false, RatingFormat::DxRatingPayload) => { let data = DxCalculatedEntries::from_user_rating_lossy(&rating.user_rating); let payload = DxRatingNet { calculated_entries: data, version: DataVersion::Prism, region: "_generic", }; json_display(payload)?; } } } commands::Commands::Logout { user_id, timestamp } => { let logout: UserLogoutApiResp = Sdgb1_50::request( &client, APIMethod::UserLogoutApi, user_id, UserLogoutApi { user_id, date_time: timestamp, ..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?; human_readable_display(preview, human_readable)?; } commands::Commands::Ping => { let time = SystemTime::now(); let decoded: PingResp = Sdgb1_50::request(&client, APIMethod::Ping, "", Ping {}).await?; info!( "sdgb 1.50 resp: {decoded}, {}ms", time.elapsed().unwrap_or_default().as_millis() ); } commands::Commands::QRLogin { ref qrcode_content } => { let qrcode = QRCode { qrcode_content }; let resp = qrcode.login(&client).await; match &resp { Ok(user_id) => info!("login succeed: {user_id}"), Err(e) => error!("login failed: {e}"), } if !human_readable { json_display(resp)?; } } commands::Commands::AuthLite { title_ver, variant } => { let resp = match variant { commands::AuthLiteVariant::SDGB => delivery_raw::(&client, title_ver).await?, commands::AuthLiteVariant::SDHJ => delivery_raw::(&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>::Ok(resp?) }) .buffer_unordered(concurrency) // slower to avoid being banned .filter_map(async |r| r.ok()) .collect::>() .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 file = OpenOptions::new() .create(true) .truncate(true) .write(true) .open("players.json")?; #[cfg(file_lock_ready)] file.try_lock()?; let txn = read_txn()?; let table = open_table_read(&txn)?; let config = bincode::config::Configuration::::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::>(); 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`");