Unify Repair + Add Policy Duplicate Detect + Confirmation Before Destructive Changes (#2)

Reviewed-on: #2
Co-authored-by: Quaternions <krakow20@gmail.com>
Co-committed-by: Quaternions <krakow20@gmail.com>

src/main.rs

@@ -1,6 +1,7 @@
 use clap::{Args,Parser,Subcommand};
 use futures::{StreamExt,TryStreamExt};
 use rand::seq::SliceRandom;
+use submissions_api::types::{ScriptResponse,ScriptPolicyResponse};
 use std::io::Write;
 use std::path::PathBuf;
 
@@ -18,8 +19,7 @@ struct Cli{
 #[derive(Subcommand)]
 enum Commands{
 	Release(ReleaseCommand),
-	RepairDuplicates(RepairDuplicatesCommand),
-	RepairPolicies(RepairPoliciesCommand),
+	Repair(RepairCommand),
 	Review(ReviewCommand),
 	UploadScripts(UploadScriptsCommand),
 }
@@ -32,14 +32,7 @@ struct ReleaseCommand{
 	api_url:String,
 }
 #[derive(Args)]
-struct RepairDuplicatesCommand{
-	#[arg(long)]
-	session_id_file:PathBuf,
-	#[arg(long)]
-	api_url:String,
-}
-#[derive(Args)]
-struct RepairPoliciesCommand{
+struct RepairCommand{
 	#[arg(long)]
 	session_id_file:PathBuf,
 	#[arg(long)]
@@ -68,11 +61,7 @@ async fn main(){
 			session_id:std::fs::read_to_string(command.session_id_file).unwrap(),
 			api_url:command.api_url,
 		}).await.unwrap(),
-		Commands::RepairDuplicates(command)=>repair_duplicates(RepairDuplicatesConfig{
-			session_id:std::fs::read_to_string(command.session_id_file).unwrap(),
-			api_url:command.api_url,
-		}).await.unwrap(),
-		Commands::RepairPolicies(command)=>repair_policies(RepairPoliciesConfig{
+		Commands::Repair(command)=>repair(RepairConfig{
 			session_id:std::fs::read_to_string(command.session_id_file).unwrap(),
 			api_url:command.api_url,
 		}).await.unwrap(),
@@ -475,92 +464,27 @@ async fn upload_scripts(config:UploadConfig)->Result<(),ScriptUploadError>{
 
 #[allow(dead_code)]
 #[derive(Debug)]
-enum RepairPoliciesError{
+enum RepairError{
+	Io(std::io::Error),
 	Cookie(submissions_api::CookieError),
 	Reqwest(submissions_api::ReqwestError),
 	GetPolicies(submissions_api::Error),
-	GetScripts(submissions_api::types::ScriptSingleItemError),
+	GetScripts(submissions_api::types::Error),
 	UpdateScriptPolicy(submissions_api::Error),
-}
-
-struct RepairPoliciesConfig{
-	session_id:String,
-	api_url:String,
-}
-async fn repair_policies(config:RepairPoliciesConfig)->Result<(),RepairPoliciesError>{
-	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(RepairPoliciesError::Cookie)?;
-	let api=&submissions_api::external::Context::new(config.api_url,cookie).map_err(RepairPoliciesError::Reqwest)?;
-
-	const LIMIT:u32=100;
-	let mut page=1;
-	loop{
-		println!("Downloading page {page}...");
-		let policies=api.get_script_policies(submissions_api::types::GetScriptPoliciesRequest{
-			Page:page,
-			Limit:LIMIT,
-			FromScriptHash:None,
-			ToScriptID:None,
-			Policy:Some(submissions_api::types::Policy::Replace),
-		}).await.map_err(RepairPoliciesError::GetPolicies)?;
-
-		futures::stream::iter(policies.iter().map(Ok)).try_for_each_concurrent(REMOTE_CONCURRENCY,async|policy|{
-			let from_script=api.get_script_from_hash(submissions_api::types::HashRequest{
-				hash:policy.FromScriptHash.as_str(),
-			}).await.map_err(RepairPoliciesError::GetScripts)?;
-
-			if let Some(from_script)=from_script{
-				if policy.ToScriptID==from_script.ID{
-					// invalid policy.  Reset the policy to None
-					api.update_script_policy(submissions_api::types::UpdateScriptPolicyRequest{
-						ID:policy.ID,
-						FromScriptID:None,
-						ToScriptID:None,
-						Policy:Some(submissions_api::types::Policy::None),
-					}).await.map_err(RepairPoliciesError::UpdateScriptPolicy)?;
-					println!("Policy updated! {:?}",policy.ID);
-				}
-			}else{
-				println!("Script did not exist! hash={}",policy.FromScriptHash);
-			}
-			Ok(())
-		}).await?;
-
-		if policies.len()<LIMIT as usize{
-			// We scanned all policies
-			println!("Done!");
-			break;
-		}else{
-			page+=1;
-		}
-	}
-
-	Ok(())
-}
-
-#[allow(dead_code)]
-#[derive(Debug)]
-enum RepairDuplicatesError{
-	Cookie(submissions_api::CookieError),
-	Reqwest(submissions_api::ReqwestError),
-	GetScripts(submissions_api::Error),
 	DeleteScript(submissions_api::Error),
+	DeleteScriptPolicy(submissions_api::Error),
 }
-
-struct RepairDuplicatesConfig{
+struct RepairConfig{
 	session_id:String,
 	api_url:String,
 }
-async fn repair_duplicates(config:RepairDuplicatesConfig)->Result<(),RepairDuplicatesError>{
-	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(RepairDuplicatesError::Cookie)?;
-	let api=&submissions_api::external::Context::new(config.api_url,cookie).map_err(RepairDuplicatesError::Reqwest)?;
-
-	let mut sources=std::collections::HashSet::new();
-
+async fn download_scripts(api:&submissions_api::external::Context)->Result<Vec<ScriptResponse>,RepairError>{
+	let mut scripts=Vec::new();
 	const LIMIT:u32=100;
 	let mut page=1;
 	loop{
-		println!("Downloading page {page}...");
-		let scripts=api.get_scripts(submissions_api::types::GetScriptsRequest{
+		println!("Downloading scripts page {page}...");
+		let new_scripts=api.get_scripts(submissions_api::types::GetScriptsRequest{
 			Page:page,
 			Limit:LIMIT,
 			Name:None,
@@ -568,27 +492,180 @@ async fn repair_duplicates(config:RepairDuplicatesConfig)->Result<(),RepairDupli
 			Source:None,
 			ResourceType:None,
 			ResourceID:None,
-		}).await.map_err(RepairDuplicatesError::GetScripts)?;
+		}).await.map_err(RepairError::GetScripts)?;
 
-		let done=scripts.len()<LIMIT as usize;
-
-		for script in scripts{
-			if !sources.insert(script.Source){
-				println!("Deleting duplicate script {:?}",script.ID);
-				api.delete_script(submissions_api::types::GetScriptRequest{
-					ScriptID:script.ID,
-				}).await.map_err(RepairDuplicatesError::DeleteScript)?;
-			}
-		}
+		let done=new_scripts.len()<LIMIT as usize;
+		scripts.extend(new_scripts);
 
 		if done{
-			// We scanned all policies
-			println!("Done!");
+			// We scanned all scripts
+			println!("Downloaded all scripts!");
 			break;
 		}else{
 			page+=1;
 		}
 	}
+	Ok(scripts)
+}
+async fn download_policies(api:&submissions_api::external::Context)->Result<Vec<ScriptPolicyResponse>,RepairError>{
+	let mut policies=Vec::new();
+	const LIMIT:u32=100;
+	let mut page=1;
+	loop{
+		println!("Downloading policies page {page}...");
+		let new_policies=api.get_script_policies(submissions_api::types::GetScriptPoliciesRequest{
+			Page:page,
+			Limit:LIMIT,
+			FromScriptHash:None,
+			ToScriptID:None,
+			Policy:Some(submissions_api::types::Policy::Replace),
+		}).await.map_err(RepairError::GetPolicies)?;
+
+		let done=new_policies.len()<LIMIT as usize;
+		policies.extend(new_policies);
+
+		if done{
+			// We scanned all policies
+			println!("Downloaded all policies!");
+			break;
+		}else{
+			page+=1;
+		}
+	}
+	Ok(policies)
+}
+async fn repair(config:RepairConfig)->Result<(),RepairError>{
+	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(RepairError::Cookie)?;
+	let api=&submissions_api::external::Context::new(config.api_url,cookie).map_err(RepairError::Reqwest)?;
+
+	let (scripts,policies)=tokio::try_join!(
+		download_scripts(api),
+		download_policies(api),
+	)?;
+
+	let mut script_id_from_hash=std::collections::HashMap::new();
+	let mut unique_sources=std::collections::HashSet::new();
+
+	let mut duplicate_scripts=Vec::new();
+
+	for script in &scripts{
+		// if not unique
+		if !unique_sources.insert(script.Source.as_str()){
+			println!("Identified duplicate script {:?}",script.ID);
+			duplicate_scripts.push(submissions_api::types::GetScriptRequest{
+				ScriptID:script.ID,
+			});
+		}else{
+			script_id_from_hash.insert(script.Hash.as_str(),script);
+		}
+	}
+
+	if !duplicate_scripts.is_empty(){
+		// ask to confirm delete scripts
+		print!("Delete {} duplicate scripts? [y/N]: ",duplicate_scripts.len());
+		std::io::stdout().flush().map_err(RepairError::Io)?;
+
+		let mut input=String::new();
+		std::io::stdin().read_line(&mut input).map_err(RepairError::Io)?;
+		match input.trim(){
+			"y"|"Y"=>(),
+			_=>{
+				println!("Quitting.");
+				return Ok(());
+			},
+		}
+		futures::stream::iter(duplicate_scripts).map(Ok).try_for_each_concurrent(REMOTE_CONCURRENCY,|request|{
+			api.delete_script(request)
+		}).await.map_err(RepairError::DeleteScript)?;
+	}
+
+	let mut unique_policies=std::collections::HashSet::new();
+
+	let mut update_policies=Vec::new();
+	let mut policies_not_unique=Vec::new();
+	let mut policies_missing_from_script=Vec::new();
+
+	for policy in &policies{
+		let from_script=script_id_from_hash.get(policy.FromScriptHash.as_str());
+
+		if let Some(&from_script)=from_script{
+			if policy.ToScriptID==from_script.ID{
+				// invalid policy.  Reset the policy to None
+				println!("Invalid policy {:?}, queueing update...",policy.ID);
+				update_policies.push(submissions_api::types::UpdateScriptPolicyRequest{
+					ID:policy.ID,
+					FromScriptID:None,
+					ToScriptID:None,
+					Policy:Some(submissions_api::types::Policy::None),
+				});
+			}else{
+				// if not unique
+				if !unique_policies.insert(policy.FromScriptHash.as_str()){
+					println!("Policy is not unique! hash={} id={:?}",policy.FromScriptHash,policy.ID);
+					policies_not_unique.push(submissions_api::types::GetScriptPolicyRequest{
+						ScriptPolicyID:policy.ID,
+					});
+				}
+			}
+		}else{
+			println!("FromScript does not exist! hash={} id={:?}",policy.FromScriptHash,policy.ID);
+			policies_missing_from_script.push(submissions_api::types::GetScriptPolicyRequest{
+				ScriptPolicyID:policy.ID,
+			});
+		}
+	}
+
+	if !update_policies.is_empty(){
+		print!("Update {} policies? [y/N]: ",update_policies.len());
+		std::io::stdout().flush().map_err(RepairError::Io)?;
+
+		let mut input=String::new();
+		std::io::stdin().read_line(&mut input).map_err(RepairError::Io)?;
+		match input.trim(){
+			"y"|"Y"=>(),
+			_=>{
+				println!("Quitting.");
+				return Ok(());
+			},
+		}
+		futures::stream::iter(update_policies).map(Ok).try_for_each_concurrent(REMOTE_CONCURRENCY,|request|{
+			api.update_script_policy(request)
+		}).await.map_err(RepairError::UpdateScriptPolicy)?;
+	}
+	if !policies_not_unique.is_empty(){
+		print!("Delete {} duplicate policies? [y/N]: ",policies_not_unique.len());
+		std::io::stdout().flush().map_err(RepairError::Io)?;
+
+		let mut input=String::new();
+		std::io::stdin().read_line(&mut input).map_err(RepairError::Io)?;
+		match input.trim(){
+			"y"|"Y"=>(),
+			_=>{
+				println!("Quitting.");
+				return Ok(());
+			},
+		}
+		futures::stream::iter(policies_not_unique).map(Ok).try_for_each_concurrent(REMOTE_CONCURRENCY,|request|{
+			api.delete_script_policy(request)
+		}).await.map_err(RepairError::DeleteScriptPolicy)?;
+	}
+	if !policies_missing_from_script.is_empty(){
+		print!("Delete {} orphaned policies? [y/N]: ",policies_missing_from_script.len());
+		std::io::stdout().flush().map_err(RepairError::Io)?;
+
+		let mut input=String::new();
+		std::io::stdin().read_line(&mut input).map_err(RepairError::Io)?;
+		match input.trim(){
+			"y"|"Y"=>(),
+			_=>{
+				println!("Quitting.");
+				return Ok(());
+			},
+		}
+		futures::stream::iter(policies_missing_from_script).map(Ok).try_for_each_concurrent(REMOTE_CONCURRENCY,|request|{
+			api.delete_script_policy(request)
+		}).await.map_err(RepairError::DeleteScriptPolicy)?;
+	}
 
 	Ok(())
 }