From b434dce0f6f0cc2ac4a9a20726aed0767366a4b1 Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Tue, 12 Sep 2023 16:28:10 -0700
Subject: [PATCH] helper funcs

---
 src/main.rs | 84 +++++++++++++++++++++++++++++++++++------------------
 1 file changed, 55 insertions(+), 29 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index a8aa286..aa86a6f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -78,6 +78,54 @@ fn get_scripts(dom:rbx_dom_weak::WeakDom) -> Vec<rbx_dom_weak::Instance>{
     scripts
 }
 
+fn get_id() -> Result<u32, Box<dyn std::error::Error>>{
+    match std::fs::read_to_string("id"){
+        Ok(id_file)=>Ok(id_file.parse::<u32>()?),
+        Err(e) => match e.kind() {
+            std::io::ErrorKind::NotFound => Ok(0),//implicitly take on id=0
+            _ => Err(e)?,
+        }
+    }
+}
+
+fn get_set_from_file(file:&str) -> Result<std::collections::HashSet<String>, Box<dyn std::error::Error>>{
+    let mut set=std::collections::HashSet::<String>::new();
+    for entry in std::fs::read_dir(file)? {
+        set.insert(std::fs::read_to_string(entry?.path())?);
+    }
+    Ok(set)
+}
+
+fn get_allowed_set() -> Result<std::collections::HashSet<String>, Box<dyn std::error::Error>>{
+    get_set_from_file("scripts/allowed")
+}
+
+fn get_blocked() -> Result<std::collections::HashSet<String>, Box<dyn std::error::Error>>{
+    get_set_from_file("scripts/blocked")
+}
+
+fn get_allowed_map() -> Result<std::collections::HashMap::<u32,String>, Box<dyn std::error::Error>>{
+    let mut allowed_map = std::collections::HashMap::<u32,String>::new();
+    for entry in std::fs::read_dir("scripts/allowed")? {
+        let entry=entry?;
+        allowed_map.insert(entry.file_name().to_str().unwrap().parse::<u32>()?,std::fs::read_to_string(entry.path())?);
+    }
+    Ok(allowed_map)
+}
+
+fn get_replace_map() -> Result<std::collections::HashMap::<String,u32>, Box<dyn std::error::Error>>{
+    let mut replace = std::collections::HashMap::<String,u32>::new();
+    for entry in std::fs::read_dir("scripts/replace")? {
+        let entry=entry?;
+        replace.insert(std::fs::read_to_string(entry.path())?,entry.file_name().to_str().unwrap().parse::<u32>()?);
+    }
+    Ok(replace)
+}
+
+fn check_source_illegal_keywords(source:&String)->bool{
+    source.find("getfenv").is_some()||source.find("require").is_some()
+}
+
 fn download(map_list: Vec<u64>) -> Result<(), Box<dyn std::error::Error>>{
     let header=format!("Cookie: .ROBLOSECURITY={}",std::env::var("RBXCOOKIE")?);
     let shared_args=&[
@@ -102,23 +150,10 @@ enum Scan{
     Flagged,
 }
 fn scan() -> Result<(), Box<dyn std::error::Error>>{
-    let mut id = 0u32;
-    match std::fs::read_to_string("id"){
-        Ok(id_file)=>id=id_file.parse::<u32>()?,
-        Err(e) => match e.kind() {
-            std::io::ErrorKind::NotFound => println!("id file does not exist: starting from 0"),//continue on, implicitly take on id=0, write the id file at the end
-            _ => return Err(e)?,
-        }
-    }
+    let mut id = get_id()?;
     //Construct allowed scripts
-    let mut allowed_set = std::collections::HashSet::<String>::new();
-    for entry in std::fs::read_dir("scripts/allowed")? {
-        allowed_set.insert(std::fs::read_to_string(entry?.path())?);
-    }
-    let mut blocked = std::collections::HashSet::<String>::new();
-    for entry in std::fs::read_dir("scripts/blocked")? {
-        blocked.insert(std::fs::read_to_string(entry?.path())?);
-    }
+    let allowed_set = get_allowed_set()?;
+    let mut blocked = get_blocked()?;
 
     for entry in std::fs::read_dir("maps/unprocessed")? {
         let file_thing=entry?;
@@ -134,7 +169,7 @@ fn scan() -> Result<(), Box<dyn std::error::Error>>{
         for script in scripts.iter() {
             if let Some(rbx_dom_weak::types::Variant::String(s)) = script.properties.get("Source") {
                 //flag keywords and instantly fail
-                if s.find("getfenv").is_some()||s.find("require").is_some(){
+                if check_source_illegal_keywords(s){
                     println!("{:?} - flagged.",file_thing.file_name());
                     fail_type=Scan::Flagged;
                     break;
@@ -169,17 +204,8 @@ fn scan() -> Result<(), Box<dyn std::error::Error>>{
     Ok(())
 }
 fn replace() -> Result<(), Box<dyn std::error::Error>>{
-    //Construct allowed scripts
-    let mut allowed_map = std::collections::HashMap::<u32,String>::new();
-    for entry in std::fs::read_dir("scripts/allowed")? {
-        let entry=entry?;
-        allowed_map.insert(entry.file_name().to_str().unwrap().parse::<u32>()?,std::fs::read_to_string(entry.path())?);
-    }
-    let mut replace = std::collections::HashMap::<String,u32>::new();
-    for entry in std::fs::read_dir("scripts/replace")? {
-        let entry=entry?;
-        replace.insert(std::fs::read_to_string(entry.path())?,entry.file_name().to_str().unwrap().parse::<u32>()?);
-    }
+    let allowed_map=get_allowed_map()?;
+    let replace_map=get_replace_map()?;
 
     for entry in std::fs::read_dir("maps/purgatory")? {
         let file_thing=entry?;
@@ -196,7 +222,7 @@ fn replace() -> Result<(), Box<dyn std::error::Error>>{
         let mut any_failed=false;
         for script in scripts.iter() {
             if let Some(rbx_dom_weak::types::Variant::String(source)) = script.properties.get("Source") {
-                if let (Some(replace_id),Some(replace_script))=(replace.get(source),write_dom.get_by_ref_mut(script.referent())) {
+                if let (Some(replace_id),Some(replace_script))=(replace_map.get(source),write_dom.get_by_ref_mut(script.referent())) {
                     println!("replace {}",replace_id);
                     //replace the source
                     if let Some(replace_source)=allowed_map.get(replace_id){