mapfixer/src/main.rs

use std::unimplemented;

use clap::{Args, Parser, Subcommand};

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

#[derive(Subcommand)]
enum Commands {
    Download(MapList),
    Upload,
    Scan,
    Replace,
    Interactive,
}

#[derive(Args)]
struct MapList {
    maps: Vec<u64>,
}

fn class_is_a(class: &str, superclass: &str) -> bool {
    if class==superclass {
        return true
    }
    let class_descriptor=rbx_reflection_database::get().classes.get(class);
    if let Some(descriptor) = &class_descriptor {
        if let Some(class_super) = &descriptor.superclass {
            return class_is_a(&class_super, superclass)
        }
    }
    return false
}

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();
    while let Some(parent)=dom.get_by_ref(pref){
        full_name.insert(0, '.');
        full_name.insert_str(0, &parent.name);
        pref=parent.parent();
    }
    full_name
}

//download
    //download list of maps to maps/unprocessed
//scan (scripts)
    //iter maps/unprocessed
    //passing moves to maps/verified
    //failing moves to maps/blocked
//replace (edits & deletions)
    //iter maps/blocked
    //replace scripts and put in maps/unprocessed
//upload
    //iter maps/verified
    //interactively print DisplayName/Creator and ask for target upload ids
//interactive
    //iter maps/unprocessed
    //for each unique script, load it into the file current.lua and have it open in sublime text
    //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<rbx_dom_weak::Instance>{
    let mut scripts = std::vec::Vec::<rbx_dom_weak::Instance>::new();

    let (_,mut instances) = dom.into_raw();
    for (_,instance) in instances.drain() {
        if class_is_a(instance.class.as_str(), "LuaSourceContainer") {
            scripts.push(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=&[
        "-q",
        "--header",
        header.as_str(),
        "-O",
    ];
    for map_id in map_list.iter() {
        std::process::Command::new("wget")
        .args(shared_args)
        .arg(format!("maps/unprocessed/{}.rbxl",map_id))
        .arg(format!("https://assetdelivery.roblox.com/v1/asset/?ID={}",map_id))
        .spawn()?;
    }
    Ok(())
}

enum Scan{
    Passed,
    Blocked,
    Flagged,
}
fn scan() -> Result<(), Box<dyn std::error::Error>>{
    let mut id = get_id()?;
    //Construct allowed scripts
    let allowed_set = get_allowed_set()?;
    let mut blocked = get_blocked()?;

    for entry in std::fs::read_dir("maps/unprocessed")? {
        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 scripts = get_scripts(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;
                    }
                }
            }else{
                panic!("FATAL: failed to get source for {:?}",file_thing.file_name());
            }
        }
        let mut dest=match fail_type {
            Scan::Passed => std::path::PathBuf::from("maps/processed"),
            Scan::Blocked => {
                println!("{:?} - {} {} not allowed.",file_thing.file_name(),fail_count,if fail_count==1 {"script"}else{"scripts"});
                std::path::PathBuf::from("maps/blocked")
            }
            Scan::Flagged => std::path::PathBuf::from("maps/flagged")
        };
        dest.push(file_thing.file_name());
        std::fs::rename(file_thing.path(), dest)?;
    }
    std::fs::write("id",id.to_string())?;
    Ok(())
}
fn replace() -> Result<(), Box<dyn std::error::Error>>{
    let allowed_map=get_allowed_map()?;
    let replace_map=get_replace_map()?;

    for entry in std::fs::read_dir("maps/blocked")? {
        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 write_dom = rbx_dom_weak::WeakDom::new(rbx_dom_weak::InstanceBuilder::empty());
        dom.clone_into_external(dom.root_ref(), &mut write_dom);

        let scripts = get_scripts(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()));
                    }else{
                        println!("failed to get replacement source {}",replace_id);
                        any_failed=true;
                    }
                }else{
                    println!("failed to failed to get replace_id and replace_script");
                    any_failed=true;
                }
            }else{
                println!("failed to failed to get source");
                any_failed=true;
            }
        }
        if any_failed {
            println!("One or more scripts failed to replace.");
        }else{
            let mut dest=std::path::PathBuf::from("maps/unprocessed");
            dest.set_file_name(file_thing.file_name());
            dest.set_extension("rbxl");//extension is always rbxl even if source file is extensionless
            let output = std::io::BufWriter::new(std::fs::File::open(dest)?);
            rbx_binary::to_writer(output, &write_dom, &[write_dom.root_ref()])?;
        }
    }
    Ok(())
}
fn upload() -> Result<(), Box<dyn std::error::Error>>{
    //interactive prompt per upload:
    //Creator: [auto fill creator]
    //DisplayName: [auto fill DisplayName]
    //id: ["New" for blank because of my double enter key]
    // std::process::Command::new("rbxcompiler")
    // .arg("--compile=false")
    // .arg("--group=6980477")
    // .arg("--asset=5692139100")
    // .arg("--input=map.rbxm")
    // .spawn()?;
    unimplemented!()
}

enum Interactive{
    Passed,
    Blocked,
    Flagged,
}
enum ScriptAction {
    Pass,
    Replace(u32),
    Flag,
    Block,
}
enum ScriptActionParseResult {
    Pass,
    Block,
}
struct ParseScriptActionErr;
impl std::str::FromStr for ScriptActionParseResult {
    type Err=ParseScriptActionErr;
    fn from_str(s: &str) -> Result<Self, Self::Err>{
        if s=="pass"||s=="1"{
            Ok(Self::Pass)
        }else if s=="block"{
            Ok(Self::Block)
        }else{
            Err(ParseScriptActionErr)
        }
    }
}

fn interactive() -> Result<(), Box<dyn std::error::Error>>{
    let mut id=get_id()?;
    //Construct allowed scripts
    let mut allowed_set=get_allowed_set()?;
    let mut allowed_map=get_allowed_map()?;
    let mut replace_map=get_replace_map()?;
    let mut blocked = get_blocked()?;

    for entry in std::fs::read_dir("maps/unprocessed")? {
        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 write_dom = rbx_dom_weak::WeakDom::new(rbx_dom_weak::InstanceBuilder::empty());
        dom.clone_into_external(dom.root_ref(), &mut write_dom);

        let scripts = get_scripts(dom);

        //check scribb
        let mut script_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::<ScriptActionParseResult>(){
                            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
                            }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;
                            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())){
                            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!("script source property is none!");
            }
        }
        let mut dest=match fail_type{
            Interactive::Passed => {
                println!("map={:?} passed with {} {}",file_thing.file_name(),script_count,if script_count==1 {"script"}else{"scripts"});
                std::path::PathBuf::from("maps/processed")
            },//write map into maps/processed
            Interactive::Blocked => {
                println!("map={:?} blocked with {}/{} {} blocked",file_thing.file_name(),block_count,script_count,if script_count==1 {"script"}else{"scripts"});
                std::path::PathBuf::from("maps/blocked")
            },//write map into maps/blocked
            Interactive::Flagged => {
                println!("map={:?} flagged",file_thing.file_name());
                std::path::PathBuf::from("maps/flagged")
            },//write map into maps/flagged
        };
        dest.push(file_thing.file_name());
        std::fs::rename(file_thing.path(), dest)?;
    }
    std::fs::write("id",id.to_string())?;
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    match cli.command {
        Commands::Download(map_list)=>download(map_list.maps),
        Commands::Upload=>upload(),
        Commands::Scan=>scan(),
        Commands::Replace=>replace(),
        Commands::Interactive=>interactive(),
    }
}