remote-script-review/src/main.rs

use clap::{Args,Parser,Subcommand};
use futures::{StreamExt,TryStreamExt};

const READ_CONCURRENCY:usize=16;
const REMOTE_CONCURRENCY:usize=16;

#[derive(Parser)]
#[command(author,version,about,long_about=None)]
#[command(propagate_version=true)]
struct Cli{
	#[command(subcommand)]
	command:Commands,
}

#[derive(Subcommand)]
enum Commands{
	Review(ReviewCommand),
	UploadScripts(UploadScriptsCommand),
}

#[derive(Args)]
struct ReviewCommand{
	#[arg(long)]
	session_id:String,
	#[arg(long)]
	api_url:String,
}
#[derive(Args)]
struct UploadScriptsCommand{
	#[arg(long)]
	session_id:String,
	#[arg(long)]
	api_url:String,
}

#[tokio::main]
async fn main(){
	let cli=Cli::parse();
	match cli.command{
		Commands::Review(command)=>review(ReviewConfig{
			session_id:command.session_id,
			api_url:command.api_url,
		}).await.unwrap(),
		Commands::UploadScripts(command)=>upload_scripts(UploadConfig{
			session_id:command.session_id,
			api_url:command.api_url,
		}).await.unwrap(),
	}
}

enum ScriptActionParseResult{
	Pass,
	Block,
	Exit,
	Delete,
}
struct ParseScriptActionErr;
impl std::str::FromStr for ScriptActionParseResult{
	type Err=ParseScriptActionErr;
	fn from_str(s:&str)->Result<Self,Self::Err>{
		match s{
			"pass\n"=>Ok(Self::Pass),
			"block\n"=>Ok(Self::Block),
			"exit\n"=>Ok(Self::Exit),
			"delete\n"=>Ok(Self::Delete),
			_=>Err(ParseScriptActionErr),
		}
	}
}

#[allow(dead_code)]
#[derive(Debug)]
enum ReviewError{
	Cookie(submissions_api::CookieError),
	Reqwest(submissions_api::ReqwestError),
	GetPolicies(submissions_api::Error),
	GetScriptFromHash(submissions_api::types::SingleItemError),
	NoScript,
	WriteCurrent(std::io::Error),
	ActionIO(std::io::Error),
	ReadCurrent(std::io::Error),
	DeduplicateModified(submissions_api::types::SingleItemError),
	UploadModified(submissions_api::Error),
	UpdateScriptPolicy(submissions_api::Error),
}

struct ReviewConfig{
	session_id:String,
	api_url:String,
}

async fn review(config:ReviewConfig)->Result<(),ReviewError>{
	// download unreviewed policies
	// review them
	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(ReviewError::Cookie)?;
	let api=submissions_api::external::Context::new(config.api_url,cookie).map_err(ReviewError::Reqwest)?;

	let unreviewed_policies=api.get_script_policies(submissions_api::types::GetScriptPoliciesRequest{
		Page:1,
		Limit:100,
		FromScriptHash:None,
		ToScriptID:None,
		Policy:Some(submissions_api::types::Policy::None),
	}).await.map_err(ReviewError::GetPolicies)?;

	for unreviewed_policy in unreviewed_policies{
		// download source code
		let script_response=api.get_script_from_hash(submissions_api::types::HashRequest{
			hash:unreviewed_policy.FromScriptHash.as_str(),
		}).await
		.map_err(ReviewError::GetScriptFromHash)?
		.ok_or(ReviewError::NoScript)?;

		let source=script_response.Source;
		//load source into current.lua
		tokio::fs::write("current.lua",source.as_str()).await.map_err(ReviewError::WriteCurrent)?;

		//prompt action in terminal
		//wait for input
		let script_action;
		loop{
			print!("action: ");
			std::io::Write::flush(&mut std::io::stdout()).map_err(ReviewError::ActionIO)?;
			let mut action_string=String::new();
			std::io::stdin().read_line(&mut action_string).map_err(ReviewError::ActionIO)?;
			if let Ok(parsed_script_action)=action_string.parse::<ScriptActionParseResult>(){
				script_action=parsed_script_action;
				break;
			}
		}
		// default to_script_id is from from_script_id (only changed for replace policy)
		let mut to_script_id=None;
		// interpret action
		let reviewed_policy=match script_action{
			ScriptActionParseResult::Pass=>{
				//if current.lua was updated, create an allowed and replace file and set script_action to replace(new_id)
				let modified_source=tokio::fs::read_to_string("current.lua").await.map_err(ReviewError::ReadCurrent)?;
				if modified_source==source{
					submissions_api::types::Policy::Allowed
				}else{
					// compute hash
					let mut hasher=siphasher::sip::SipHasher::new();
					std::hash::Hasher::write(&mut hasher,source.as_bytes());
					let hash=std::hash::Hasher::finish(&hasher);

					// check if modified script already exists
					let maybe_script_response=api.get_script_from_hash(submissions_api::types::HashRequest{
						hash:format!("{:016x}",hash).as_str(),
					}).await.map_err(ReviewError::DeduplicateModified)?;

					// write to_script_id, uploading modified script if necessary
					to_script_id=Some(match maybe_script_response{
						Some(script_response)=>script_response.ID,
						None=>api.create_script(submissions_api::types::CreateScriptRequest{
							Name:script_response.Name.as_str(),
							Source:modified_source.as_str(),
							SubmissionID:Some(script_response.SubmissionID),
						}).await.map_err(ReviewError::UploadModified)?.ID
					});

					// use replace policy
					submissions_api::types::Policy::Replace
				}
			},
			ScriptActionParseResult::Block=>submissions_api::types::Policy::Blocked,
			ScriptActionParseResult::Exit=>break,
			ScriptActionParseResult::Delete=>submissions_api::types::Policy::Delete,
		};

		// update policy
		api.update_script_policy(submissions_api::types::UpdateScriptPolicyRequest{
			ScriptPolicyID:unreviewed_policy.ID,
			FromScriptID:None,
			ToScriptID:to_script_id,
			Policy:Some(reviewed_policy),
		}).await.map_err(ReviewError::UpdateScriptPolicy)?;
	}

	Ok(())
}

#[allow(dead_code)]
#[derive(Debug)]
enum ScriptUploadError{
	Cookie(submissions_api::CookieError),
	Reqwest(submissions_api::ReqwestError),
	AllowedSet(std::io::Error),
	AllowedMap(GetMapError),
	ReplaceMap(GetMapError),
	BlockedSet(std::io::Error),
	GOC(GOCError),
	GOCPolicyReplace(GOCError),
	GOCPolicyAllowed(GOCError),
	GOCPolicyBlocked(GOCError),
}

fn read_dir_stream(dir:tokio::fs::ReadDir)->impl futures::stream::Stream<Item=std::io::Result<tokio::fs::DirEntry>>{
    futures::stream::unfold(dir,|mut dir|async{
        match dir.next_entry().await{
            Ok(Some(entry))=>Some((Ok(entry),dir)),
            Ok(None)=>None, // End of directory
            Err(e)=>Some((Err(e),dir)), // Error encountered
        }
    })
}

async fn get_set_from_file(path:impl AsRef<std::path::Path>)->std::io::Result<std::collections::HashSet<String>>{
	read_dir_stream(tokio::fs::read_dir(path).await?)
	.map(|dir_entry|async{
		tokio::fs::read_to_string(dir_entry?.path()).await
	})
	.buffer_unordered(READ_CONCURRENCY)
	.try_collect().await
}

async fn get_allowed_set()->std::io::Result<std::collections::HashSet<String>>{
	get_set_from_file("scripts/allowed").await
}

async fn get_blocked_set()->std::io::Result<std::collections::HashSet<String>>{
	get_set_from_file("scripts/blocked").await
}

#[allow(dead_code)]
#[derive(Debug)]
enum GetMapError{
	IO(std::io::Error),
	FileStem,
	ToStr,
	ParseInt(std::num::ParseIntError),
}

async fn get_allowed_map()->Result<std::collections::HashMap::<u32,String>,GetMapError>{
	read_dir_stream(tokio::fs::read_dir("scripts/allowed").await.map_err(GetMapError::IO)?)
	.map(|dir_entry|async{
		let path=dir_entry.map_err(GetMapError::IO)?.path();
		let id:u32=path.file_stem()
		.ok_or(GetMapError::FileStem)?
		.to_str()
		.ok_or(GetMapError::ToStr)?
		.parse().map_err(GetMapError::ParseInt)?;
		let source=tokio::fs::read_to_string(path).await.map_err(GetMapError::IO)?;
		Ok((id,source))
	})
	.buffer_unordered(READ_CONCURRENCY)
	.try_collect().await
}

async fn get_replace_map()->Result<std::collections::HashMap::<String,u32>,GetMapError>{
	read_dir_stream(tokio::fs::read_dir("scripts/replace").await.map_err(GetMapError::IO)?)
	.map(|dir_entry|async{
		let path=dir_entry.map_err(GetMapError::IO)?.path();
		let id:u32=path.file_stem()
		.ok_or(GetMapError::FileStem)?
		.to_str()
		.ok_or(GetMapError::ToStr)?
		.parse().map_err(GetMapError::ParseInt)?;
		let source=tokio::fs::read_to_string(path).await.map_err(GetMapError::IO)?;
		Ok((source,id))
	})
	.buffer_unordered(READ_CONCURRENCY)
	.try_collect().await
}

fn hash_source(source:&str)->u64{
	let mut hasher=siphasher::sip::SipHasher::new();
	std::hash::Hasher::write(&mut hasher,source.as_bytes());
	std::hash::Hasher::finish(&hasher)
}

fn hash_format(hash:u64)->String{
	format!("{:016x}",hash)
}

type GOCError=submissions_api::types::SingleItemError;
type GOCResult=Result<submissions_api::types::ScriptID,GOCError>;

async fn get_or_create_script(api:&submissions_api::external::Context,source:&str)->GOCResult{
	let script_response=api.get_script_from_hash(submissions_api::types::HashRequest{
		hash:hash_format(hash_source(source)).as_str(),
	}).await?;

	Ok(match script_response{
		Some(script_response)=>script_response.ID,
		None=>api.create_script(submissions_api::types::CreateScriptRequest{
			Name:"Script",
			Source:source,
			SubmissionID:None,
		}).await.map_err(GOCError::Other)?.ID
	})
}

async fn check_or_create_script_poicy(
	api:&submissions_api::external::Context,
	hash:&str,
	script_policy:submissions_api::types::CreateScriptPolicyRequest,
)->Result<(),GOCError>{
	let script_policy_result=api.get_script_policy_from_hash(submissions_api::types::HashRequest{
		hash,
	}).await?;

	match script_policy_result{
		Some(script_policy_reponse)=>{
			// check that everything matches the expectation
			assert!(hash==script_policy_reponse.FromScriptHash);
			assert!(script_policy.ToScriptID==script_policy_reponse.ToScriptID);
			assert!(script_policy.Policy==script_policy_reponse.Policy);
		},
		None=>{
			// create a new policy
			api.create_script_policy(script_policy).await.map_err(GOCError::Other)?;
		}
	}

	Ok(())
}

async fn do_policy(
	api:&submissions_api::external::Context,
	script_ids:&std::collections::HashMap<&str,submissions_api::types::ScriptID>,
	source:&str,
	to_script_id:submissions_api::types::ScriptID,
	policy:submissions_api::types::Policy,
)->Result<(),GOCError>{
	let hash=hash_format(hash_source(source));
	check_or_create_script_poicy(api,hash.as_str(),submissions_api::types::CreateScriptPolicyRequest{
		FromScriptID:script_ids[source],
		ToScriptID:to_script_id,
		Policy:policy,
	}).await
}

struct UploadConfig{
	session_id:String,
	api_url:String,
}

async fn upload_scripts(config:UploadConfig)->Result<(),ScriptUploadError>{
	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(ScriptUploadError::Cookie)?;
	let api=&submissions_api::external::Context::new(config.api_url,cookie).map_err(ScriptUploadError::Reqwest)?;

	let (
		allowed_set_result,
		allowed_map_result,
		replace_map_result,
		blocked_set_result,
	)=tokio::join!(
		get_allowed_set(),
		get_allowed_map(),
		get_replace_map(),
		get_blocked_set(),
	);

	let allowed_set=allowed_set_result.map_err(ScriptUploadError::AllowedSet)?;
	let allowed_map=allowed_map_result.map_err(ScriptUploadError::AllowedMap)?;
	let replace_map=replace_map_result.map_err(ScriptUploadError::ReplaceMap)?;
	let blocked_set=blocked_set_result.map_err(ScriptUploadError::BlockedSet)?;

	// create a unified deduplicated set of all scripts
	let script_set:std::collections::HashSet<&str>=allowed_set.iter()
	.map(String::as_str)
	.chain(
		replace_map.keys().map(String::as_str)
	).chain(
		blocked_set.iter().map(String::as_str)
	).collect();

	// get or create every unique script
	let script_ids:std::collections::HashMap<&str,submissions_api::types::ScriptID>=
	futures::stream::iter(script_set)
	.map(|source|async move{
		let script_id=get_or_create_script(api,source).await?;
		Ok((source,script_id))
	})
	.buffer_unordered(REMOTE_CONCURRENCY)
	.try_collect().await.map_err(ScriptUploadError::GOC)?;

	// get or create policy for each script in each category
	//
	// replace
	let replace_fut=futures::stream::iter(replace_map.iter().map(Ok))
	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|(source,id)|async{
		do_policy(
			api,
			&script_ids,
			source,
			script_ids[allowed_map[id].as_str()],
			submissions_api::types::Policy::Replace
		).await.map_err(ScriptUploadError::GOCPolicyReplace)
	});

	// allowed
	let allowed_fut=futures::stream::iter(allowed_set.iter().map(Ok))
	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|source|async{
		do_policy(
			api,
			&script_ids,
			source,
			script_ids[source.as_str()],
			submissions_api::types::Policy::Allowed
		).await.map_err(ScriptUploadError::GOCPolicyAllowed)
	});

	// blocked
	let blocked_fut=futures::stream::iter(blocked_set.iter().map(Ok))
	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|source|async{
		do_policy(
			api,
			&script_ids,
			source,
			script_ids[source.as_str()],
			submissions_api::types::Policy::Blocked
		).await.map_err(ScriptUploadError::GOCPolicyBlocked)
	});

	tokio::try_join!(replace_fut,allowed_fut,blocked_fut)?;

	Ok(())
}