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;
type AssetIDFileMap=Vec<(AssetID,std::path::PathBuf)>;
const CONCURRENT_REQUESTS:usize=8;

/// Parse a single key-value pair
fn parse_key_val<T,U>(s:&str)->AResult<(T,U)>
where
	T:std::str::FromStr,
	T::Err:std::error::Error+Send+Sync+'static,
	U:std::str::FromStr,
	U::Err:std::error::Error+Send+Sync+'static,
{
	let pos=s
		.find('=')
		.ok_or_else(||anyhow::Error::msg(format!("invalid KEY=value: no `=` found in `{s}`")))?;
	Ok((s[..pos].parse()?,s[pos+1..].parse()?))
}

#[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>,

	#[arg(long,value_parser=parse_key_val::<AssetID,std::path::PathBuf>)]
	asset_ids:AssetIDFileMap,

	#[command(subcommand)]
	command:Commands,
}

#[derive(Subcommand)]
enum Commands{
	Download,
	Upload,
}

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

#[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=>download_list(cookie,cli.asset_ids).await,
		Commands::Upload=>upload_list(cookie,group,cli.asset_ids).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,asset_id_file_map:AssetIDFileMap)->AResult<()>{
	let client=reqwest::Client::new();
	futures::stream::iter(asset_id_file_map)
	.map(|(asset_id,file)|{
		let client=&client;
		let cookie=cookie.as_str();
		let owner=&owner;
		async move{
			let mut url=reqwest::Url::parse("https://data.roblox.com/Data/Upload.ashx?json=1&type=Model&genreTypeId=1")?;
			//url borrow scope
			{
				let mut query=url.query_pairs_mut();//borrow here
				query.append_pair("assetid",asset_id.to_string().as_str());
				match owner{
					Owner::Group(group_id)=>{query.append_pair("groupId",group_id.to_string().as_str());},
					Owner::User=>(),
				}
			}
			
			let body=tokio::fs::read_to_string(file).await?;
			let mut resp=client.post(url.clone())
			.header("Cookie",cookie)
			.body(body.clone())
			.send().await?;

			//This is called a CSRF challenge apparently
			if resp.status()==reqwest::StatusCode::FORBIDDEN{
				if let Some(csrf_token)=resp.headers().get("X-CSRF-Token"){
					resp=client.post(url)
					.header("X-CSRF-Token",csrf_token)
					.header("Cookie",cookie)
					.body(body)
					.send().await?;
				}else{
					return Err(anyhow::Error::msg("Roblox returned 403 with no CSRF"));
				}
			}

			Ok((asset_id,resp.bytes().await?))
		}
	})
	.buffer_unordered(CONCURRENT_REQUESTS)
	.for_each(|b:AResult<_>|async{
			match b{
				Ok((asset_id,body))=>{
					println!("asset_id={} response.body={:?}",asset_id,body);
				},
				Err(e)=>eprintln!("ul error: {}",e),
			}
		}).await;
	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_id_file_map:AssetIDFileMap)->AResult<()>{
	let client=reqwest::Client::new();
	futures::stream::iter(asset_id_file_map)
	.map(|(asset_id,file)|{
		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((file,resp.bytes().await?))
		}
	})
	.buffer_unordered(CONCURRENT_REQUESTS)
	.for_each(|b:AResult<_>|async{
			match b{
				Ok((dest,body))=>{
					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(())
}