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(AssetIDFileMapBad),
	Upload(AssetIDFileMapBad),
}

//idk how to make this a list of key-value pairs
#[derive(Args)]
struct AssetIDFileMapBad{
	asset_ids:Vec<AssetID>,
	files:Vec<std::path::PathBuf>,
}

#[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(asset_id_file_map)=>download_list(cookie,transpose_asset_id_file_map(asset_id_file_map)?).await,
		Commands::Upload(asset_id_file_map)=>upload_list(cookie,group,transpose_asset_id_file_map(asset_id_file_map)?).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"))
	}
}

type AssetIDFileMap=Vec<(AssetID,std::path::PathBuf)>;

fn transpose_asset_id_file_map(asset_id_file_map:AssetIDFileMapBad)->AResult<AssetIDFileMap>{
	if asset_id_file_map.asset_ids.len()==asset_id_file_map.files.len(){
		Ok(asset_id_file_map.asset_ids.into_iter().zip(asset_id_file_map.files.into_iter()).collect())
	}else{
		Err(anyhow::Error::msg("Asset list did not match file list."))
	}
}

async fn upload_list(cookie:String,owner:Owner,asset_id_file_map:AssetIDFileMap)->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_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(())
}