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;
const CONCURRENT_REQUESTS:usize=8;

#[derive(Parser)]
#[command(author,version,about,long_about=None)]
#[command(propagate_version = true)]
struct Cli{
	#[arg(short,long)]
	group:Option<u64>,
	//idk how to do this better
	#[arg(long)]
	cookie_literal:Option<String>,
	#[arg(long)]
	cookie_env:Option<String>,
	#[arg(long)]
	cookie_file:Option<std::path::PathBuf>,

	#[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();

	let cookie_enum={
		match (cli.cookie_literal,cli.cookie_env,cli.cookie_file){
			(Some(literal),None,None)=>Cookie::Literal(literal),
			(None,Some(env_var),None)=>Cookie::Environment(env_var),
			(None,None,Some(path))=>Cookie::File(path),
			_=>return Err(anyhow::Error::msg("Cookie was not specified or was specified multiple times.")),
		}
	};
	let cookie=format!(".ROBLOSECURITY={}",match cookie_enum{
		Cookie::Literal(s)=>s,
		Cookie::Environment(var)=>std::env::var(var)?,
		Cookie::File(path)=>tokio::fs::read_to_string(path).await?,
	});

	let group=match cli.group{
		Some(group_id)=>Owner::Group(group_id),
		None=>Owner::User,
	};

	match cli.command{
		Commands::Download(asset_id_list)=>download_list(cookie,asset_id_list.asset_ids).await,
		Commands::Upload{path,asset_id}=>upload_list(cookie,group,path,asset_id).await,
	}
}

enum Owner{
	Group(u64),
	User
}

enum Cookie{
	Literal(String),
	Environment(String),
	File(std::path::PathBuf),
}

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"))
	}
}

async fn upload_list(_cookie:String,_owner:Owner,_path:std::path::PathBuf,_asset_id:AssetID)->AResult<()>{
	Ok(())
}

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(cookie:String,asset_ids:Vec<AssetID>)->AResult<()>{
	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(())
}