From 1075d59a916a5425a3593d43576ae6b9cba1fec7 Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Sun, 31 Dec 2023 11:15:27 -0800
Subject: [PATCH] async upload

---
 src/main.rs | 92 +++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 71 insertions(+), 21 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 761c197..e36e55f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,8 +4,23 @@ 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)]
@@ -20,21 +35,17 @@ struct Cli{
 	#[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(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>,
+	Download,
+	Upload,
 }
 
 #[derive(Args)]
@@ -66,8 +77,8 @@ async fn main()->AResult<()>{
 	};
 
 	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,
+		Commands::Download=>download_list(cookie,cli.asset_ids).await,
+		Commands::Upload=>upload_list(cookie,group,cli.asset_ids).await,
 	}
 }
 
@@ -99,17 +110,56 @@ fn maybe_gzip_decode<R:Read+Seek>(input:&mut R)->AResult<ReaderType<R>>{
 	}
 }
 
-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<()>{
+	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(())
 }