command_injector/src/main.rs

use std::str::from_utf8_unchecked;
use std::{str::from_utf8, process::Command};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write, Seek, SeekFrom};
use std::ffi::OsString;
use std::path::Path;
#[derive(Debug)]
struct Ops {
    folder: Option<String>,
    output: Option<String>,
    recursive: bool,
    
}

// Father please forgive me for the sins im about to commit
static mut HOME_DIR: String = String::new();

fn main() {
    let ops = read_ops();
    println!("{:?}", ops);
    if ops.folder.is_none() || ops.output.is_none() {
        panic!("missing folder/output!");
    }
    unsafe { HOME_DIR = ops.folder.clone().unwrap(); }
    inject_folder(&ops.folder.unwrap(), &ops.output.unwrap(), ops.recursive);
}

fn inject_folder(folder: &str, output: &str, rec: bool) {
    // Create the directory just in case
    let _ = fs::create_dir_all(output);
    
    let f = fs::read_dir(folder).expect(&format!("Couldn't open folder {}", folder));
    for i in f.filter(|x| x.is_ok()).map(|x| x.unwrap()) {
        if i.file_name().to_str().unwrap().starts_with('.') {
            continue; // Invis folder/file, we dont want to copy them
        }
        let pbuf = i.path();
        let p = pbuf.as_path();
        if p.is_file() {
            let op = String::from_iter([output, "/", i.file_name().to_str().unwrap()]);
            inject_file(&p, &Path::new(&op));
        }
        else if rec && p.is_dir() {
            let fol = p.as_os_str().to_str().unwrap();
            let out = String::from_iter([output, "/", i.file_name().to_str().unwrap()]);
            inject_folder(fol, &out, rec);
        }
    }
}

fn inject_file(file: &Path, output: &Path) {
    println!("injecting {:?} to {:?}", file, output);
    let mut inp = OpenOptions::new().read(true).open(file).expect(&format!("Cannot open {:?}", file.as_os_str()));
    let mut out = File::create(output).expect(&format!("Cannot create file {:?}", output.as_os_str()));
    let mut buf: [u8; 2048] = [0; 2048];


    if file.extension().unwrap_or(&OsString::new()) == "html" {
        // Copy while also reading and searching for <INJECT> ... </INJECT>
        while let Ok(len) = inp.read(&mut buf) {
            if len == 0 {
                // EOF
                break;
            }
            match from_utf8(&buf[..len]) {
                Ok(s) => {
                    write_injection(
                        s, 
                        &mut out, 
                        &mut inp, 
                        file.parent().unwrap().as_os_str().to_str().unwrap(),
                        file.file_name().unwrap().to_str().unwrap(),
                    );
                },
                Err(e) => {
                    if e.valid_up_to() == 0 {
                        panic!("File is not valid UTF8");
                    }
                    // Safety: we know all bytes are valid from the error...
                    let s = unsafe { from_utf8_unchecked(&buf[..e.valid_up_to()]) };
                    inp.seek(SeekFrom::Current(-1 * (len - e.valid_up_to()) as i64)).unwrap();

                    write_injection(
                        s, 
                        &mut out, 
                        &mut inp, 
                        file.parent().unwrap().as_os_str().to_str().unwrap(),
                        file.file_name().unwrap().to_str().unwrap(),
                    );
                },
            }

        }
    }
    else {
        // Just copy the file!
        while let Ok(len) = inp.read(&mut buf) {
            if len == 0 {
                // EOF
                break;
            }
            let res = out.write(&buf[..len]);
            if res.is_err() {
                panic!("Cannot write to {:?}", output.as_os_str());
            }
            else if res.unwrap() != len {
                panic!("cannot continue writing to {:?}", output.as_os_str());
            }
        }
    }
    out.flush().expect(&format!("Cannot flush file {:?}", output.as_os_str()));
}

fn write_injection(s: &str, out: &mut File, inp: &mut File, cfolder: &str, cfile: &str) {
    let tag : &str = "<INJECT>";
    let ctag : &str = "</INJECT>";

    if s.contains(tag) && s.contains(ctag) {
        let (pre,post) = s.split_once(tag).unwrap();
        let (com, post) = post.split_once(ctag).unwrap();

        out.write(pre.as_bytes()).unwrap();

        // let mut com_iter = com.split(' ');
        let com = Command::new("bash")
            .arg("-c")
            .arg(com)
            .env("HOME", unsafe { &HOME_DIR })
            .env("CURRENT", cfolder)
            .env("FILE", cfile)
            .output();

        let com = match com {
            Ok(o) => String::from_utf8(o.stdout).unwrap_or(String::from("UTF8 conversion error")),
            Err(_) => String::from("Failed running command :("),
        };

        out.write(com.as_bytes()).unwrap();
        write_injection(post, out, inp, cfolder, cfile);
    }
    else if s.contains(tag) {
        // We have the first tag but not the second, so we want to go back in the input file..
        let (copy, keep) = s.split_once(tag).unwrap();

        out.write(copy.as_bytes()).unwrap();
        // Return the length of whatever we just read and the tag length
        inp.seek(SeekFrom::Current(-1 * (keep.as_bytes().len() + tag.as_bytes().len()) as i64)).unwrap();
    }
    else {
        // Make sure we are not just cutting it in the middle
        let chars = tag.chars().collect::<Vec<_>>();
        let mut wrote = false;
        for i in 1..chars.len() {
            if s.ends_with(&chars[..i]) {
                let back: usize = chars[..i].iter().map(|x| x.len_utf8()).sum();
                let to_write = s.replace(&chars[..i], "");
                out.write(to_write.as_bytes()).unwrap();

                inp.seek(SeekFrom::Current(-1 * back as i64)).unwrap();
                wrote = true;
                break;
            }
        }
        if !wrote {
            out.write(s.as_bytes()).unwrap();
        }
    }
}

fn read_ops() -> Ops {
    let args = std::env::args();
    let mut ops = Ops { folder: None, output: None, recursive: false };

    for a in args.skip(1) {
        if a.starts_with('-') {
            if a.trim() == "-r" || a.trim() == "--recursive" {
                ops.recursive = true;
            }
        }
        else if ops.folder.is_none() {
            ops.folder = Some(a);
        }
        else if ops.output.is_none() {
            ops.output = Some(a);
        }
    }

    ops
}