asset-tool/src/main.rs

use std::io::{Read,Seek};
use clap::{Args,Parser,Subcommand};
use anyhow::Result as AResult;
use futures::StreamExt;

type AssetID=u64;

#[derive(Parser)]
#[command(author,version,about,long_about=None)]
#[command(propagate_version = true)]
struct Cli{
	#[command(subcommand)]
	command:Commands,
}

#[derive(Subcommand)]
enum Commands{
	Download(AssetIDList),
	Upload{path:std::path::PathBuf,asset_id:AssetID},
}

#[derive(Args)]
struct PathBufList{
	paths:Vec<std::path::PathBuf>
}

#[derive(Args)]
struct AssetIDList{
	asset_ids:Vec<AssetID>,
}

#[tokio::main]
async fn main()->AResult<()>{
	let cli=Cli::parse();
	match cli.command{
		Commands::Download(asset_id_list)=>download_list(asset_id_list.asset_ids).await,
		Commands::Upload{path,asset_id}=>upload_file(path,asset_id),
	}
}

enum ReaderType<'a,R:Read+Seek>{
	GZip(flate2::read::GzDecoder<&'a mut R>),
	Raw(&'a mut R),
}

fn maybe_gzip_decode<R:Read+Seek>(input:&mut R)->AResult<ReaderType<R>>{
	let mut first_2=[0u8;2];
	if let (Ok(()),Ok(()))=(std::io::Read::read_exact(input,&mut first_2),std::io::Seek::rewind(input)){
		match &first_2{
			b"\x1f\x8b"=>Ok(ReaderType::GZip(flate2::read::GzDecoder::new(input))),
			_=>Ok(ReaderType::Raw(input)),
		}
	}else{
		Err(anyhow::Error::msg("failed to peek"))
	}
}

fn upload_file(_path:std::path::PathBuf,_asset_id:AssetID)->AResult<()>{
	Ok(())
}

const CONCURRENT_REQUESTS:usize=8;

fn read_readable(mut readable:impl Read)->AResult<Vec<u8>>{
	let mut contents=Vec::new();
	readable.read_to_end(&mut contents)?;
	Ok(contents)
}

async fn download_list(asset_ids:Vec<AssetID>)->AResult<()>{
	let cookie=format!(".ROBLOSECURITY={}",std::env::var("RBXCOOKIE_PROJECTSLIME")?);
	let client=reqwest::Client::new();
	futures::stream::iter(asset_ids)
	.map(|asset_id|{
		let client=&client;
		let cookie=cookie.as_str();
		async move{
			let resp=client.get(format!("https://assetdelivery.roblox.com/v1/asset/?ID={}",asset_id))
			.header("Cookie",cookie)
			.send().await?;
			Ok((asset_id,resp.bytes().await?))
		}
	})
	.buffer_unordered(CONCURRENT_REQUESTS)
	.for_each(|b:AResult<_>|async{
			match b{
				Ok((asset_id,body))=>{
					let dest=std::path::PathBuf::from(asset_id.to_string());
					let contents=match maybe_gzip_decode(&mut std::io::Cursor::new(body)){
						Ok(ReaderType::GZip(readable))=>read_readable(readable),
						Ok(ReaderType::Raw(readable))=>read_readable(readable),
						Err(e)=>Err(e),
					};
					match contents{
						Ok(data)=>match tokio::fs::write(dest,data).await{
							Err(e)=>eprintln!("fs error: {}",e),
							_=>(),
						},
						Err(e)=>eprintln!("gzip error: {}",e),
					};
				},
				Err(e)=>eprintln!("dl error: {}",e),
			}
		}).await;
	Ok(())
}