async upload

This commit is contained in:
Quaternions 2023-12-31 11:15:27 -08:00
parent 01676bb4da
commit 1075d59a91

View File

@ -4,8 +4,23 @@ use anyhow::Result as AResult;
use futures::StreamExt; use futures::StreamExt;
type AssetID=u64; type AssetID=u64;
type AssetIDFileMap=Vec<(AssetID,std::path::PathBuf)>;
const CONCURRENT_REQUESTS:usize=8; 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)] #[derive(Parser)]
#[command(author,version,about,long_about=None)] #[command(author,version,about,long_about=None)]
#[command(propagate_version = true)] #[command(propagate_version = true)]
@ -20,21 +35,17 @@ struct Cli{
#[arg(long)] #[arg(long)]
cookie_file:Option<std::path::PathBuf>, cookie_file:Option<std::path::PathBuf>,
#[arg(long,value_parser=parse_key_val::<AssetID,std::path::PathBuf>)]
asset_ids:AssetIDFileMap,
#[command(subcommand)] #[command(subcommand)]
command:Commands, command:Commands,
} }
#[derive(Subcommand)] #[derive(Subcommand)]
enum Commands{ enum Commands{
Download(AssetIDFileMapBad), Download,
Upload(AssetIDFileMapBad), Upload,
}
//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)] #[derive(Args)]
@ -66,8 +77,8 @@ async fn main()->AResult<()>{
}; };
match cli.command{ match cli.command{
Commands::Download(asset_id_file_map)=>download_list(cookie,transpose_asset_id_file_map(asset_id_file_map)?).await, Commands::Download=>download_list(cookie,cli.asset_ids).await,
Commands::Upload(asset_id_file_map)=>upload_list(cookie,group,transpose_asset_id_file_map(asset_id_file_map)?).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<()>{ 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(()) Ok(())
} }