diff --git a/src/main.rs b/src/main.rs index 20da11f..42a554f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -42,7 +42,16 @@ fn class_is_a(class: &str, superclass: &str) -> bool { } return false } - +fn recursive_collect_scripts(scripts: &mut std::vec::Vec,dom: &rbx_dom_weak::WeakDom, instance: &rbx_dom_weak::Instance){ + for &referent in instance.children() { + if let Some(c) = dom.get_by_ref(referent) { + if class_is_a(c.class.as_str(), "LuaSourceContainer") { + scripts.push(c.referent());//copy ref + } + recursive_collect_scripts(scripts,dom,c); + } + } +} fn get_full_name(dom:&rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance) -> String{ let mut full_name=instance.name.clone(); let mut pref=instance.parent(); @@ -72,15 +81,9 @@ fn get_full_name(dom:&rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance) -> //I can edit the file and it will edit it in place //I pass/fail(with comment)/allow each script -fn get_scripts(dom:rbx_dom_weak::WeakDom) -> Vec{ - let mut scripts = std::vec::Vec::::new(); - - let (_,mut instances) = dom.into_raw(); - for (_,instance) in instances.drain() { - if class_is_a(instance.class.as_str(), "LuaSourceContainer") { - scripts.push(instance); - } - } +fn get_script_refs(dom:&rbx_dom_weak::WeakDom) -> Vec{ + let mut scripts = std::vec::Vec::new(); + recursive_collect_scripts(&mut scripts, dom, dom.root()); scripts } @@ -167,32 +170,36 @@ fn scan() -> Result<(), Box>{ let dom = rbx_binary::from_reader(input)?; - let scripts = get_scripts(dom); + let script_refs = get_script_refs(&dom); //check scribb let mut fail_count=0; let mut fail_type=Scan::Passed; - 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 check_source_illegal_keywords(s){ - println!("{:?} - flagged.",file_thing.file_name()); - fail_type=Scan::Flagged; - break; - } - if allowed_set.contains(s) { - continue; - }else{ - fail_type=Scan::Blocked;//no need to check for Flagged, it breaks the loop. - fail_count+=1; - if !blocked.contains(s) { - blocked.insert(s.clone());//all fixed! just clone! - std::fs::write(format!("scripts/blocked/{}.lua",id),s)?; - id+=1; + for &script_ref in script_refs.iter() { + if let Some(script)=dom.get_by_ref(script_ref){ + if let Some(rbx_dom_weak::types::Variant::String(s)) = script.properties.get("Source") { + //flag keywords and instantly fail + if check_source_illegal_keywords(s){ + println!("{:?} - flagged.",file_thing.file_name()); + fail_type=Scan::Flagged; + break; } + if allowed_set.contains(s) { + continue; + }else{ + fail_type=Scan::Blocked;//no need to check for Flagged, it breaks the loop. + fail_count+=1; + if !blocked.contains(s) { + blocked.insert(s.clone());//all fixed! just clone! + std::fs::write(format!("scripts/blocked/{}.lua",id),s)?; + id+=1; + } + } + }else{ + panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); } }else{ - panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); + panic!("FATAL: failed to get_by_ref {:?}",script_ref); } } let mut dest=match fail_type { @@ -225,20 +232,24 @@ fn extract(file_id:u64) -> Result<(), Box>{ let dom = rbx_binary::from_reader(input)?; - let scripts = get_scripts(dom); + let script_refs = get_script_refs(&dom); //extract scribb - for script in scripts.iter() { - if let Some(rbx_dom_weak::types::Variant::String(s)) = script.properties.get("Source") { - if script_set.contains(s) { - continue; + for &script_ref in script_refs.iter() { + if let Some(script)=dom.get_by_ref(script_ref){ + if let Some(rbx_dom_weak::types::Variant::String(s)) = script.properties.get("Source") { + if script_set.contains(s) { + continue; + }else{ + script_set.insert(s.clone()); + std::fs::write(format!("scripts/extracted/{:?}_{}_{}.lua",file_thing.file_name(),id,script.name),s)?; + id+=1; + } }else{ - script_set.insert(s.clone()); - std::fs::write(format!("scripts/extracted/{:?}_{}_{}.lua",file_thing.file_name(),id,script.name),s)?; - id+=1; + panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); } }else{ - panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); + panic!("FATAL: failed to get_by_ref {:?}",script_ref); } } } @@ -253,33 +264,33 @@ fn replace() -> Result<(), Box>{ let file_thing=entry?; let input = std::io::BufReader::new(std::fs::File::open(file_thing.path())?); - let dom = rbx_binary::from_reader(input)?; + let mut dom = rbx_binary::from_reader(input)?; - let input = std::io::BufReader::new(std::fs::File::open(file_thing.path())?); - let mut write_dom = rbx_binary::from_reader(input)?; - - let scripts = get_scripts(dom); + let script_refs = get_script_refs(&dom); //check scribb 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_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){ - replace_script.properties.insert("Source".to_string(), rbx_dom_weak::types::Variant::String(replace_source.clone())); + for &script_ref in script_refs.iter() { + if let Some(script)=dom.get_by_ref(script_ref){ + if let Some(rbx_dom_weak::types::Variant::String(source)) = script.properties.get("Source") { + if let (Some(replace_id),Some(replace_script))=(replace_map.get(source),dom.get_by_ref_mut(script.referent())) { + println!("replace {}",replace_id); + //replace the source + if let Some(replace_source)=allowed_map.get(replace_id){ + replace_script.properties.insert("Source".to_string(), rbx_dom_weak::types::Variant::String(replace_source.clone())); + }else{ + println!("failed to get replacement source {}",replace_id); + any_failed=true; + } }else{ - println!("failed to get replacement source {}",replace_id); + println!("failed to failed to get replace_id and replace_script"); any_failed=true; } }else{ - println!("failed to failed to get replace_id and replace_script"); - any_failed=true; + panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); } }else{ - println!("failed to failed to get source"); - any_failed=true; + panic!("FATAL: failed to get_by_ref {:?}",script_ref); } } if any_failed { @@ -288,7 +299,7 @@ fn replace() -> Result<(), Box>{ let mut dest=std::path::PathBuf::from("maps/unprocessed"); dest.set_file_name(file_thing.file_name()); let output = std::io::BufWriter::new(std::fs::File::open(dest)?); - rbx_binary::to_writer(output, &write_dom, &[write_dom.root_ref()])?; + rbx_binary::to_writer(output, &dom, &[dom.root_ref()])?; } } Ok(()) @@ -348,115 +359,116 @@ fn interactive() -> Result<(), Box>{ let file_thing=entry?; println!("processing map={:?}",file_thing.file_name()); let input = std::io::BufReader::new(std::fs::File::open(file_thing.path())?); - let dom = rbx_binary::from_reader(input)?; + let mut dom = rbx_binary::from_reader(input)?; - let input = std::io::BufReader::new(std::fs::File::open(file_thing.path())?); - let mut write_dom = rbx_binary::from_reader(input)?; - - let scripts = get_scripts(dom); + let script_refs = get_script_refs(&dom); //check scribb let mut script_count=0; let mut replace_count=0; let mut block_count=0; let mut fail_type=Interactive::Passed; - for script in scripts.iter() { - if let Some(rbx_dom_weak::types::Variant::String(source)) = script.properties.get("Source") { - script_count+=1; - let source_action=if check_source_illegal_keywords(source) { - ScriptAction::Flag//script triggers flagging -> Flag - } else if blocked.contains(source) { - ScriptAction::Block//script is blocked -> Block - } else if allowed_set.contains(source) { - ScriptAction::Pass//script is allowed -> Pass - }else if let Some(replace_id)=replace_map.get(source) { - ScriptAction::Replace(*replace_id) - }else{ - //interactive logic goes here - //load source into current.lua - std::fs::write("current.lua",source)?; - //prompt action in terminal - println!("unresolved source location={}",get_full_name(&write_dom, script)); - //wait for input - let script_action; - loop{ - let mut action_string = String::new(); - std::io::stdin().read_line(&mut action_string)?; - if let Ok(parsed_script_action)=action_string.parse::(){ - script_action=parsed_script_action; - break; - }else{ - println!("illegal action string."); - } - } - //update allowed/replace/blocked - 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=std::fs::read_to_string("current.lua")?; - if &modified_source==source{ - //it's always new. - //insert allowed_set - allowed_set.insert(modified_source.clone()); - //insert allowed_map - allowed_map.insert(id,modified_source.clone()); - //write allowed/id.lua - std::fs::write(format!("scripts/allowed/{}.lua",id),modified_source)?; - id+=1; - ScriptAction::Pass + for &script_ref in script_refs.iter() { + if let Some(script)=dom.get_by_ref(script_ref){ + if let Some(rbx_dom_weak::types::Variant::String(source)) = script.properties.get("Source") { + script_count+=1; + let source_action=if check_source_illegal_keywords(source) { + ScriptAction::Flag//script triggers flagging -> Flag + } else if blocked.contains(source) { + ScriptAction::Block//script is blocked -> Block + } else if allowed_set.contains(source) { + ScriptAction::Pass//script is allowed -> Pass + }else if let Some(replace_id)=replace_map.get(source) { + ScriptAction::Replace(*replace_id) + }else{ + //interactive logic goes here + //load source into current.lua + std::fs::write("current.lua",source)?; + //prompt action in terminal + println!("unresolved source location={}",get_full_name(&dom, script)); + //wait for input + let script_action; + loop{ + let mut action_string = String::new(); + std::io::stdin().read_line(&mut action_string)?; + if let Ok(parsed_script_action)=action_string.parse::(){ + script_action=parsed_script_action; + break; }else{ - //insert allowed_set - allowed_set.insert(modified_source.clone()); - //insert allowed_map - allowed_map.insert(id,modified_source.clone()); - //insert replace_map - replace_map.insert(source.clone(),id);//this cannot be reached if it already exists - //write allowed/id.lua - std::fs::write(format!("scripts/allowed/{}.lua",id),modified_source)?; - //write replace/id.lua - std::fs::write(format!("scripts/replace/{}.lua",id),source)?; - let ret=ScriptAction::Replace(id); + println!("illegal action string."); + } + } + //update allowed/replace/blocked + 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=std::fs::read_to_string("current.lua")?; + if &modified_source==source{ + //it's always new. + //insert allowed_set + allowed_set.insert(modified_source.clone()); + //insert allowed_map + allowed_map.insert(id,modified_source.clone()); + //write allowed/id.lua + std::fs::write(format!("scripts/allowed/{}.lua",id),modified_source)?; + id+=1; + ScriptAction::Pass + }else{ + //insert allowed_set + allowed_set.insert(modified_source.clone()); + //insert allowed_map + allowed_map.insert(id,modified_source.clone()); + //insert replace_map + replace_map.insert(source.clone(),id);//this cannot be reached if it already exists + //write allowed/id.lua + std::fs::write(format!("scripts/allowed/{}.lua",id),modified_source)?; + //write replace/id.lua + std::fs::write(format!("scripts/replace/{}.lua",id),source)?; + let ret=ScriptAction::Replace(id); + id+=1; + ret + } + }, + ScriptActionParseResult::Block => { + blocked.insert(source.clone()); + std::fs::write(format!("scripts/blocked/{}.lua",id),source)?; id+=1; - ret + ScriptAction::Block + }, + } + }; + + match source_action{ + ScriptAction::Pass => println!("passed source location={}",get_full_name(&dom, script)), + ScriptAction::Replace(replace_id)=>{ + //replace the source + let location=get_full_name(&dom, script); + if let (Some(replace_source),Some(replace_script))=(allowed_map.get(&replace_id),dom.get_by_ref_mut(script.referent())){ + replace_count+=1; + println!("replaced source id={} location={}",replace_id,location); + replace_script.properties.insert("Source".to_string(), rbx_dom_weak::types::Variant::String(replace_source.clone())); + }else{ + panic!("failed to get replacement source id={} location={}",replace_id,location); + } + } + ScriptAction::Flag => { + println!("flagged source location={}",get_full_name(&dom, script)); + fail_type=Interactive::Flagged; + }, + ScriptAction::Block => { + block_count+=1; + println!("blocked source location={}",get_full_name(&dom, script)); + match fail_type{ + Interactive::Passed => fail_type=Interactive::Blocked, + _=>(), } }, - ScriptActionParseResult::Block => { - blocked.insert(source.clone()); - std::fs::write(format!("scripts/blocked/{}.lua",id),source)?; - id+=1; - ScriptAction::Block - }, } - }; - - match source_action{ - ScriptAction::Pass => println!("passed source location={}",get_full_name(&write_dom, script)), - ScriptAction::Replace(replace_id)=>{ - //replace the source - let location=get_full_name(&write_dom, script); - if let (Some(replace_source),Some(replace_script))=(allowed_map.get(&replace_id),write_dom.get_by_ref_mut(script.referent())){ - replace_count+=1; - println!("replaced source id={} location={}",replace_id,location); - replace_script.properties.insert("Source".to_string(), rbx_dom_weak::types::Variant::String(replace_source.clone())); - }else{ - panic!("failed to get replacement source id={} location={}",replace_id,location); - } - } - ScriptAction::Flag => { - println!("flagged source location={}",get_full_name(&write_dom, script)); - fail_type=Interactive::Flagged; - }, - ScriptAction::Block => { - block_count+=1; - println!("blocked source location={}",get_full_name(&write_dom, script)); - match fail_type{ - Interactive::Passed => fail_type=Interactive::Blocked, - _=>(), - } - }, + }else{ + panic!("FATAL: failed to get source for {:?}",file_thing.file_name()); } }else{ - panic!("script source property is none!"); + panic!("FATAL: failed to get_by_ref {:?}",script_ref); } } let mut dest=match fail_type{