diff --git a/Cargo.toml b/Cargo.toml index ef58e3fa83..45e5e12ad5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "libafl_tinyinst", "libafl_sugar", "libafl_nyx", + "libafl_wizard", "libafl_concolic/symcc_runtime", "libafl_concolic/symcc_libafl", "libafl_concolic/test/dump_constraints", diff --git a/libafl_wizard/Cargo.toml b/libafl_wizard/Cargo.toml new file mode 100644 index 0000000000..e9f5e6b13f --- /dev/null +++ b/libafl_wizard/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "libafl_wizard" +version.workspace = true +edition = "2021" +authors = ["matheus "] +description = "libafl's fuzzer template generator" +documentation = "TODO" +repository = "https://github.com/AFLplusplus/LibAFL/" +readme = "../README.md" +license = "MIT OR Apache-2.0" +keywords = ["fuzzing", "testing", "security"] +categories = ["development-tools::testing"] + +[dependencies] +toml = "0.7.6" +serde = { version = "1.0", features = ["derive"] } +graphviz-rust = "0.6.6" \ No newline at end of file diff --git a/libafl_wizard/README.md b/libafl_wizard/README.md new file mode 100644 index 0000000000..6daeb7243c --- /dev/null +++ b/libafl_wizard/README.md @@ -0,0 +1,23 @@ +# libafl_wizard + +libafl_wizard logo + +libafl_wizard is a tool to generate fuzzers using libafl's components. By answering some questions, you can generate your own fuzzer to test programs and help with the development of more secure code, all that while learning more about libafl and how it's used! + +## Usage + +libafl_wizard has a cli interface and can be run with: + +``` +cargo run +``` + +## Have in mind that... + +The tool makes use of [graphviz](https://graphviz.org/download/) to generate an image containing the flowchart of the questions diagram, so the users can know beforehand where the answers will take them. Make sure it's installed (and added to PATH) before running the tool. + +When writing answers, the check if an input is a valid answer is such that it simply verifies if what's typed by the user has the same characters so far as the answer (check out `validate_input()`). The thing is that, e.g. if "Crash or Timeout" and "Crash" are both valid answers, if the user answers "crash", the first option will be deemed correct (even though the user wanted the second one). To avoid that, for now, one can simply reverse these answers, so "Crash" comes before "Crash or Timeout". + +## Contributing + +libafl_wizard uses the `questions.toml` TOML file to store and load the questions that will be asked during the generation process. Each question contains some fields, like the possible answers for that question and the Rust code associated to those answers. As libafl's components get updated or new ones introduced, the questions need to be updated as well. diff --git a/libafl_wizard/icons/libafl_wizard.png b/libafl_wizard/icons/libafl_wizard.png new file mode 100644 index 0000000000..c12ddbfd91 Binary files /dev/null and b/libafl_wizard/icons/libafl_wizard.png differ diff --git a/libafl_wizard/questions.toml b/libafl_wizard/questions.toml new file mode 100644 index 0000000000..34699854c3 --- /dev/null +++ b/libafl_wizard/questions.toml @@ -0,0 +1,271 @@ +[[question]] +id = "intro" +title = "libafl_wizard: a tool to generate fuzzers using Libafl's components." +content = """ +Before starting, make sure that you know your target very well. Choosing the right components to build a fuzzer depends on the details +of the implementation of your target. + +Knowing what you want before building a fuzzer can be extremely helpful when selecting the components that best fit your options and +the restrictions of the environment. + +Details such as: + *Having the source code of the target. + *Implementing a harness function to fuzz your target. + *Knowing the expected input by the target. + *Providing an initial set of inputs. + *The size of the target and the availability of memory. + *The possibility of memory corruption. + +Can be helpful during this process. +""" +skipped_by = "" +previous = "intro" + + [[question.answers]] + was_chosen = false + answer = "Next" + next = "source code" + code = "" + skip = [] + +[[question]] +id = "source code" +title = "Do you have the target's source code?" +content = """ +Having the target's source code is interesting for performance. With the source code in hands, we can instrument the target, that is, +place coverage information to help our fuzzer identify new points of execution (guide itself). + +Without instrumentation on the source code, we have to rely on third-party applications to provide this information for our fuzzer, +such as QEMU, FRIDA or Tiny Inst. +""" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "Yes" + next = "harness" + code = "" + skip = [] + + [[question.answers]] + was_chosen = false + answer = "No" + next = "END" + code = "" + skip = [] + +[[question]] +id = "harness" +title = "Can you provide a harness function?" +content = """ +A harness function bridges the gap between how the fuzzer expects input to occur and how it's delivered. Essentially, the harness will +receive the input (properly structured) and call the target/function under test with it. This allows for in-process fuzzing (when the +harness is executed inside the fuzzer process), which increases performance a lot. + +Not providing a harness makes in-process fuzzing impractical. Therefore, the only option left is to fork before executing the target, +which makes use of a shared memory region to record the coverage information. +""" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "Yes" + next = "corrupt memory" + code = "" + skip = [] + + [[question.answers]] + was_chosen = false + answer = "No" + next = "crash/timeout" + code = """ + /* Forkserver Executor */ + """ + skip = [] + +[[question]] +id = "corrupt memory" +title = "Can your target corrupt memory used by the fuzzer?" +content = """ +Under some circumstances, you may find your harness pretty unstable or your harness wreaks havoc on the global states. In this case, +you want to fork it before executing the harness runs in the child process so that it doesn't break things. +""" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "Yes" + next = "crash/timeout" + code = """ + use libafl::InProcessForkExecutor; + use libafl_bolts::{ + shmem::{unix_shmem, ShMemProvider}, + tuples::tuple_list, + }; + + let mut shmem_provider = unix_shmem::UnixShMemProvider::new().unwrap(); + + let mut executor = InProcessForkExecutor::new( + &mut harness, + tuple_list!(observer), + &mut fuzzer, + &mut state, + &mut mgr, + shmem_provider, + ) + .expect("Failed to create the Executor"); + """ + skip = [] + + [[question.answers]] + was_chosen = false + answer = "No" + next = "crash/timeout" + code = """ + use libafl::executors::inprocess::InProcessExecutor; + use libafl_bolts::tuples::tuple_list; + + let mut executor = InProcessExecutor::new( + &mut harness, + tuple_list!(observer), + &mut fuzzer, + &mut state, + &mut mgr, + ) + .expect("Failed to create the Executor"); + """ + skip = [] + +[[question]] +id = "crash/timeout" +title = "Do you expect to cause a crash or a timeout on the target?" +content = """ +Determining the objective of the fuzzing campaign is essential to identify input that triggered critical errors on the target. + +Telling the fuzzer that we are looking for a crash means that a testcase, which causes a crash on the target, fullfills the objective +and, differently from the ones that e.g. reach a new execution point, this testcase is saved in a separate folder for you to check the +input that trigerred the crash. + +A timeout follows the same idea: a testcase that takes longer to execute than a particular timeout. + +It's even possible to join these two ideas to instruct the fuzzzer that any testcase, which causes a crash or takes long enough to +execute, is stored. +""" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "Crash" + next = "default components" + code = """ + use libafl::feedbacks::CrashFeedback; + + let mut objective = CrashFeedback::new(); + """ + skip = [] + + [[question.answers]] + was_chosen = false + answer = "Crash or Timeout" + next = "default components" + code = """ + use libafl::{ + feedback_or_fast, + feedbacks::{CrashFeedback, TimeoutFeedback}, + }; + + let mut objective = feedback_or_fast!( + CrashFeedback::new(), + TimeoutFeedback::new() + ); + """ + skip = [] + + [[question.answers]] + was_chosen = false + answer = "Timeout" + next = "default components" + code = """ + use libafl::feedbacks::TimeoutFeedback; + + let mut objective = TimeoutFeedback::new(); + """ + skip = [] + +[[question]] +id = "default components" +title = "Type in finsish to end the generation." +content = """ +""" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "Finish" + next = "END" + code = """ + use std::path::PathBuf; + use libafl::{ + corpus::{InMemoryCorpus, OnDiskCorpus}, + events::SimpleEventManager, + feedbacks::MaxMapFeedback, + fuzzer::{Fuzzer, StdFuzzer}, + monitors::SimpleMonitor, + mutators::scheduled::{havoc_mutations, StdScheduledMutator}, + schedulers::QueueScheduler, + stages::mutational::StdMutationalStage, + state::StdState, + }; + use libafl_bolts::{ + current_nanos, + rands::StdRand, + tuples::tuple_list, + }; + + let mut feedback = MaxMapFeedback::new(&observer); + + let mut state = StdState::new( + StdRand::with_seed(current_nanos()), + InMemoryCorpus::new(), + OnDiskCorpus::new(PathBuf::from("./crashes")).unwrap(), + &mut feedback, + &mut objective, + ) + .unwrap(); + + let mon = SimpleMonitor::new(|s| println!("{s}")); + + let mut mgr = SimpleEventManager::new(mon); + + let scheduler = QueueScheduler::new(); + + let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective); + + let mutator = StdScheduledMutator::new(havoc_mutations()); + + let mut stages = tuple_list!(StdMutationalStage::new(mutator)); + + fuzzer + .fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr) + .expect("Error in the fuzzing loop"); + """ + skip = [] + +[[question]] +id = "END" +title = "All questions answered!" +content = "" +skipped_by = "" +previous = "" + + [[question.answers]] + was_chosen = false + answer = "" + next = "" + code = "" + skip = [] \ No newline at end of file diff --git a/libafl_wizard/src/answer.rs b/libafl_wizard/src/answer.rs new file mode 100644 index 0000000000..e1533bc29a --- /dev/null +++ b/libafl_wizard/src/answer.rs @@ -0,0 +1,55 @@ +use serde::Deserialize; + +/// The Answer struct contains all the necessary information that an answer must +/// have, such as the code to be added and the next question to be asked. +#[derive(Clone, Deserialize, PartialEq, Debug)] +pub struct Answer { + was_chosen: bool, + answer: String, + next: String, + code: String, + skip: Vec, +} + +impl Answer { + /// Returns true if this answer was chosen between all the possibilities for + /// a given question. + pub fn was_chosen(&self) -> &bool { + &self.was_chosen + } + + /// Returns true if this answer was chosen between all the possibilities for + /// a given question. + pub fn set_was_chosen(&mut self, new_value: bool) { + self.was_chosen = new_value; + } + + /// Returns a String that represents the text of that answer, e.g "Yes" or + /// "No". + pub fn answer(&self) -> &String { + &self.answer + } + + /// Returns the id of the next question that will be asked, considering that + /// this answer was chosen. + pub fn next(&self) -> &String { + &self.next + } + + /// Returns the Rust code that will be added to the fuzzer file, considering + /// that this answer was chosen. + pub fn code(&self) -> &String { + &self.code + } + + /// Returns the ids of the questions that should be skipped, considering + /// that this answer was chosen. + /// + /// In some cases, depending on the answer that the user chooses for a + /// particular question, it will skip other questions that shouldn't be + /// asked, e.g. if the user doesn't have the source code of the that target + /// the wizard shouldn't ask if they can provide a harness. + pub fn skip(&self) -> &Vec { + &self.skip + } +} diff --git a/libafl_wizard/src/main.rs b/libafl_wizard/src/main.rs new file mode 100644 index 0000000000..ca982f45bf --- /dev/null +++ b/libafl_wizard/src/main.rs @@ -0,0 +1,121 @@ +/*! +libafl_wizard is a crate that allows users to generate fuzzers that make use of +Libafl's components by simply answering some questions. + +This is the entry point into libafl_wizard. +*/ + +use std::{io, io::Write}; + +mod answer; +mod question; +mod utils; + +use question::{find_question, flowchart_image, Question}; +use utils::{ + arrange_code, arrange_imports, clear_terminal_screen, separate_imports, validate_input, + write_code, +}; + +fn main() { + // The question diagram is a vector containing all the questions. + let mut questions = Question::new(); + // Each element is a String, which contains the Rust code associated to a question. This will be used to write to the file. + let mut code_content: Vec = Vec::new(); + // Index of the current question. + let mut curr_q = 0; + // Index of the next question. Note that, when undoing, the next question is the previous one (that led to the current one). + let mut next_q; + // Index of the chosen answer. + let mut ans_i; + let mut input = String::new(); + + // Generate a flowchat image to help guide the user. + flowchart_image(&questions); + + while questions[curr_q].id() != "END" { + clear_terminal_screen(); + questions[curr_q].print_question(); + print!("\n >> "); + io::stdout().flush().unwrap(); + io::stdin() + .read_line(&mut input) + .expect("Failed to get input from stdin."); + input = input.trim().to_string(); + + while !questions[curr_q].is_answer(&input) { + print!("Please, type a valid answer: "); + io::stdout().flush().unwrap(); + input.clear(); + io::stdin() + .read_line(&mut input) + .expect("Failed to get input from stdin."); + input = input.trim().to_string(); + } + + if validate_input(&input, &String::from("Undo")) { + // The "Undo" option makes the generator go back to the previous answered question, so if the user do something by + // mistake they can correct it. + next_q = find_question(&questions, &questions[curr_q].previous()); + + // If the user chooses to undo a question that produced code, the associated code is removed. + // Since the Undo option goes backwards, we can simply pop the last piece of code. + for ans in questions[next_q].answers().iter() { + if *ans.was_chosen() && !ans.code().is_empty() { + code_content.pop(); + } + } + + // Also, if we are undoing this question and it skipped others, we undo this too. + ans_i = questions[curr_q].chosen_answer(); + let answers = questions[next_q].answers(); + + if !answers[ans_i].skip().is_empty() { + questions[next_q] + .clone() + .unskip_questions(&mut questions, ans_i); + } + let answers = questions[next_q].mut_answers(); + answers[ans_i].set_was_chosen(false); + } else { + (next_q, ans_i) = questions[curr_q].resolve_answer(&questions, &input); + let answers = questions[curr_q].answers(); + + // Adds the code associated to the user choice. + if !answers[ans_i].code().is_empty() { + code_content.push(answers[ans_i].code().to_string()); + } + + // If there are any questions that should be skipped because of that answer, we skip them. + if !answers[ans_i].skip().is_empty() { + questions[curr_q] + .clone() + .skip_questions(&mut questions, ans_i); + } + + // Only updates the 'previous' field when going forward (not undoing) in the questions diagram. + let q_id = questions[curr_q].id().clone(); + questions[next_q].set_previous(q_id); + + let answers = questions[curr_q].mut_answers(); + answers[ans_i].set_was_chosen(true); + } + input.clear(); + curr_q = next_q; + } + + let (imports_content, code_content) = separate_imports(code_content); + + println!("IMPORTS:"); + for i in &imports_content { + println!("{}", i); + } + + // Separate by instances of components/imports, arrange them in the correct order and write to the file. + let file_name = write_code(arrange_code(code_content), arrange_imports(imports_content)); + + println!( + "\nFile {} successfully created in the ./fuzzers directory.\nAll questions answered!", + file_name + ); +} diff --git a/libafl_wizard/src/question.rs b/libafl_wizard/src/question.rs new file mode 100644 index 0000000000..89070b2cb9 --- /dev/null +++ b/libafl_wizard/src/question.rs @@ -0,0 +1,305 @@ +use graphviz_rust::{ + cmd::{CommandArg, Format}, + exec, + printer::PrinterContext, +}; +use serde::Deserialize; +use std::fs::read_to_string; +use toml::from_str; + +use crate::answer::Answer; +use crate::utils::validate_input; + +/// Used to read the TOML file containing the questions. +#[derive(Deserialize)] +pub struct QuestionList { + question: Vec, +} + +/// The Question struct contains all the necessary information that a question +/// must have, such as the title containing the question being asked and the set +/// of possible answers for this particular question. +#[derive(Clone, Deserialize, PartialEq, Debug)] +pub struct Question { + id: String, + title: String, + content: String, + skipped_by: String, + previous: String, + answers: Vec, +} + +impl Question { + /// Reads the questions in the TOML file to a vector, where each field is an + /// unique question. + pub fn new() -> Vec { + let contents = read_to_string("questions.toml").expect("Failed to read questions file."); + + let q_list: QuestionList = from_str(&contents).expect("Failed to parse toml questions."); + + q_list.question + } + + /// Returns the id of the question, which is used to differentiate between + /// questions. + pub fn id(&self) -> &String { + &self.id + } + + /// Returns the title of the question, which in most cases is the questions + /// being asked. + /// + /// The 'title' is also used as the text of the nodes in the flowchart + /// image. + pub fn title(&self) -> &String { + &self.title + } + + /// Returns the description of the question, usually some information to + /// help the user understand the concepts associated with a particular + /// question. + pub fn content(&self) -> &String { + &self.content + } + + /// Returns the id of the question that skipped the one under consideration. + pub fn skipped_by(&self) -> &String { + &self.skipped_by + } + + /// Sets that this particular question was skipped by the one with this id. + pub fn set_skipped_by(&mut self, id: String) { + self.skipped_by = id; + } + + /// Returns the id of the question that led to the current one. + pub fn previous(&self) -> &String { + &self.previous + } + + /// Sets that this particular question came after the one with this id was + /// answered. + pub fn set_previous(&mut self, id: String) { + self.previous = id; + } + + /// Returns the set of possible answers for this question, excluding the + /// Undo option. + pub fn answers(&self) -> &Vec { + &self.answers + } + + /// Returns a mutable set of possible answers for this question, excluding + /// the Undo option. + pub fn mut_answers(&mut self) -> &mut Vec { + &mut self.answers + } + + /// Prints all the relevant information of this question. + pub fn print_question(&self) { + let mut output = String::new(); + + // Construct the output string + output.push_str(&format!( + "+=====================+\n| libafl_wizard |\n+=====================+\n\n" + )); + output.push_str(&format!("{}\n\n", self.title())); + output.push_str(&format!("{}\n\n\t", self.content())); + for ans in self.answers().iter() { + output.push_str(&format!( + "{}{}|{}", + ans.answer(), + " ".repeat(4), + " ".repeat(4) + )); + } + output.push_str("Undo\n"); + print!("{}", output); + } + + /// Checks if the user typed one of the acceptable answers or is undoing. + pub fn is_answer(&self, input: &String) -> bool { + if input.is_empty() { + return false; + } else if validate_input(&input, &String::from("Undo")) { + return true; + } + for ans in self.answers().iter() { + if validate_input(&input, &ans.answer()) { + return true; + } + } + + false + } + + /// Returns the index of the chosen answer in the vector of possible answer + /// for this given question. + pub fn chosen_answer(&self) -> usize { + for (i, ans) in self.answers().iter().enumerate() { + if *ans.was_chosen() { + return i; + } + } + + 0 + } + + /// Returns a tuple containing the id of the next question and the index of + /// answer chosen for this question. + /// + /// If an invalid answer is provided, the function returns 0 as default for + /// both values. + pub fn resolve_answer(&self, questions: &Vec, input: &String) -> (usize, usize) { + // Checks which of the acceptable answers the user typed. If so, returns the index of the next question associated to it. + for (i, ans) in self.answers().iter().enumerate() { + if validate_input(&input, &ans.answer()) { + let mut next_q = find_question(questions, &ans.next()); + + // If the question should be skipped, then the wizard goes to next question. + // These types of questions should always have only one possibility for next question because the wizard cant infer + // which answer the user would have chosen. + while !questions[next_q].skipped_by().is_empty() { + next_q = find_question(questions, &ans.next()); + } + + return (next_q, i); + } + } + + (0, 0) + } + + /// Marks the questions to be skipped. + pub fn skip_questions(&self, questions: &mut Vec, ans_i: usize) { + let answers = self.answers(); + + for q_id in answers[ans_i].skip().iter() { + let i = questions + .iter() + .position(|question| question.id() == q_id) + .unwrap(); + + questions[i].set_skipped_by(self.id().clone()); + } + } + + /// Unmarks the questions that would be skipped. + pub fn unskip_questions(&self, questions: &mut Vec, ans_i: usize) { + let answers = self.answers(); + + for q_id in answers[ans_i].skip().iter() { + let i = questions + .iter() + .position(|question| question.id() == q_id) + .unwrap(); + + questions[i].set_skipped_by("".to_string()); + } + } + + /// Returns true if, for the given set of answers for this question, at + /// least one leads to a different next question than the others. + pub fn diff_next_questions(&self) -> bool { + for (i, ans1) in self.answers().iter().enumerate() { + for (j, ans2) in self.answers().iter().enumerate() { + if i != j { + if ans1.next() != ans2.next() { + return true; + } + } + } + } + + false + } +} + +/// Returns the index of the question in the questions vector that has id == q. +pub fn find_question(questions: &Vec, q: &String) -> usize { + questions + .iter() + .position(|question| question.id() == q) + .unwrap() +} + +/// Generates an image containg the flowchart of the questions of the wizard. +/// +/// Requires 'graphviz' to be installed on the machine, or panics. +pub fn flowchart_image(questions: &Vec) { + let mut dot_string = String::from("digraph t {\n"); + + for q in questions { + if q.id() != "END" { + dot_string.push_str(&format!("\t\"{}\"[color=black]\n", q.title())); + + if !q.diff_next_questions() { + let j = questions + .iter() + .position(|question| &question.id() == &q.answers[0].next()) + .unwrap(); + + // Yes or No questions that lead to the same next. + if q.answers().len() <= 2 { + for ans in q.answers() { + dot_string.push_str(&format!( + "\t\"{}\" -> \"{}\"\n[label=\"{}\"]", + q.title(), + questions[j].title(), + ans.answer(), + )); + } + } else { + // Multiple answers that lead to the same next. + dot_string.push_str(&format!( + "\t\"{}\" -> \"{}\"\n[label=\"{}\"]", + q.title(), + questions[j].title(), + q.answers[0].answer(), + )); + + dot_string.push_str(&format!( + "\t\"{}\" -> \"{}\"\n[label=\"...\"]", + q.title(), + questions[j].title(), + )); + } + } + // Multiple answers that lead to distinct next questions. + else { + for ans in q.answers().iter() { + let j = questions + .iter() + .position(|question| question.id() == ans.next()) + .unwrap(); + + dot_string.push_str(&format!( + "\t\"{}\" -> \"{}\"\n[label=\"{}\"]", + q.title(), + questions[j].title(), + ans.answer(), + )); + } + } + } else { + break; + } + } + + dot_string.push_str("}"); + + let g = graphviz_rust::parse(&dot_string).unwrap(); + + let mut ctx = PrinterContext::default(); + ctx.always_inline(); + + let _graph_png = exec( + g, + &mut ctx, + vec![ + CommandArg::Format(Format::Png), + CommandArg::Output("flowchart.png".to_string()), + ], + ) + .unwrap(); +} diff --git a/libafl_wizard/src/utils.rs b/libafl_wizard/src/utils.rs new file mode 100644 index 0000000000..4804fcbece --- /dev/null +++ b/libafl_wizard/src/utils.rs @@ -0,0 +1,608 @@ +use std::{ + fs::{create_dir, OpenOptions}, + io::Write, + path::Path, + process::Command, +}; + +/// Clears the terminal screen. +pub fn clear_terminal_screen() { + if cfg!(target_os = "windows") { + Command::new("cmd") + .args(["/c", "cls"]) + .spawn() + .expect("cls command failed to start") + .wait() + .expect("failed to wait"); + } else { + Command::new("clear") + .spawn() + .expect("clear command failed to start") + .wait() + .expect("failed to wait"); + }; +} + +/// Returns true if the 'input' provided by the user is equal to 'ans'. +pub fn validate_input(input: &String, ans: &String) -> bool { + let input_low = input.to_lowercase(); + let mut input_chars = input_low.chars(); + let ans_low = ans.to_lowercase(); + let mut ans_chars = ans_low.chars(); + + // Basically, an answer is valid if it is an acceptable variant of that given answer. Acceptable variants are strings that contain + // the characters in the same order as the answer, so for the answer "Yes", acceptable variants are: "y", "Ye", "yes", "YES", but + // not "Yess", "yy", "Yhes", "yYess"... + while let Some(input_c) = input_chars.next() { + if let Some(ans_c) = ans_chars.next() { + if input_c != ans_c { + return false; + } + } else { + return false; + } + } + + true +} + +/// Returns a vector containing indivdual declarations of Libafl's components +/// and sorted according to correct order of components. +pub fn arrange_code(code_content: Vec) -> Vec { + // List of libafl's components. + let components = vec![ + "Observer", + "Feedback", + "State", + "Monitor", + "Event", + "Scheduler", + "Fuzzer", + "Executor", + "Generator", + "Mutator", + "Stage", + ]; + + let code_content = separate_code(code_content, &components); + + let mut ordered_code: Vec = Vec::new(); + let mut placed_code: Vec = vec![false; code_content.len()]; + + for comp in components { + for (i, code_line) in code_content.iter().enumerate() { + if code_line.contains(comp) { + // Place in the correct order. + ordered_code.push(code_line.to_string()); + placed_code[i] = true; + } + } + } + + // Deals with cases where there is no definition of a component. + for (i, code_line) in code_content.iter().enumerate() { + if !placed_code[i] { + ordered_code.insert(i, code_line.to_string()); + } + } + + ordered_code +} + +/// Returns a vector containing individual declarations of Libafl's components. +fn separate_code(code_content: Vec, components: &Vec<&str>) -> Vec { + let mut separated_code: Vec = Vec::new(); + + for code_string in code_content { + let mut current_line = String::new(); + let mut in_code_block = false; + + for c in code_string.chars() { + current_line.push(c); + + if !in_code_block { + if components.iter().any(|&s| current_line.contains(s)) { + in_code_block = true; + } + } + + // The end of the declaration of a component. + if in_code_block && c == ';' { + in_code_block = false; + separated_code.push(current_line.trim().to_string()); + current_line.clear(); + } + } + + if !current_line.trim().to_string().is_empty() { + separated_code.push(current_line.trim().to_string()); + } + } + + separated_code +} + +/// Creates a file and writes the imports and code to that file. +pub fn write_code(code_content: Vec, imports: Vec) -> String { + let mut counter = 0; + let mut file_name = format!("fuzzer.rs"); + + let fuzzers_folder = "./fuzzers"; + if !Path::new(fuzzers_folder).exists() { + create_dir(fuzzers_folder).expect("Failed to create fuzzers directory."); + } + + // Creates "fuzzer.rs", "fuzzer_1.rs" files if the previous one already exists... + while Path::new(&format!("{}/{}", fuzzers_folder, file_name)).exists() { + counter += 1; + file_name = format!("fuzzer_{}.rs", counter); + } + + let file_path = format!("{}/{}", fuzzers_folder, file_name); + + let mut out_file = OpenOptions::new() + .append(true) + .create(true) + .open(&file_path) + .expect("Failed to open the fuzzer file."); + + for i in imports { + out_file + .write_all(&format!("{}\n", i).as_bytes()) + .expect("Failed to write to the fuzzer file."); + } + + out_file + .write_all(&format!("\nfn main() {}", "{\n").as_bytes()) + .expect("Failed to write to the fuzzer file."); + + for code in code_content.iter() { + out_file + .write_all(&format!("{}{}\n\n", " ".repeat(4), code).as_bytes()) + .expect("Failed to write to the fuzzer file."); + } + + out_file + .write_all("}".as_bytes()) + .expect("Failed to write to the fuzzer file."); + + file_name +} + +/// Returns a tuple containing a vector with only the code of the components and +/// another containing only the imports. +pub fn separate_imports(code_content: Vec) -> (Vec, Vec) { + let mut imports: Vec = Vec::new(); + let mut new_code_content: Vec = Vec::new(); + + for code in code_content.iter() { + let mut is_import = false; + let mut import_string = String::new(); + let mut code_string = String::new(); + + for line in code.lines() { + if is_import { + import_string.push_str(&format!("{}\n", line)); + if line.contains(";") { + is_import = false; + imports.push(format_import(&import_string)); + import_string.clear(); + } + } else { + let trimmed_line = line.trim(); + if trimmed_line.starts_with("use") { + if !line.contains(";") { + is_import = true; + import_string.push_str(&format!("{}\n", line)); + } else { + imports.push(format_import(&line.to_string())); + } + } else { + if !trimmed_line.is_empty() { + code_string.push_str(&format!("{}\n", line)); + } else { + code_string.push_str("\n"); + } + } + } + } + new_code_content.push(code_string.trim_end().to_string()); + } + + (imports, new_code_content) +} + +/// Fixes identation and new line characters for the imports. +fn format_import(import: &String) -> String { + let mut new_import = String::new(); + + for (index, line) in import.lines().enumerate() { + let mut line = line.trim().to_string(); + + if index > 0 { + if line != "};" { + line.insert_str(0, " "); + line.push('\n'); + new_import.push_str(&line); + } else { + new_import.push_str(&line); + } + } else { + if !line.contains(";") { + line.push('\n'); + } + new_import.push_str(&line); + } + } + + new_import +} + +/// Arranges all the imports by sorting them alphabetically and making +/// insertions. +/// +/// For example, if there are two 'use libafl::' imports their code +/// will be joined, so that it becomes only one 'use libafl::' import. +pub fn arrange_imports(imports_content: Vec) -> Vec { + // Each field of 'result' will be a single import/use of a crate, after the insertions. + let mut result: Vec = Vec::new(); + + for import in imports_content { + let mut i: Option = None; + + // Check if the crate that we are trying to insert is already in 'result'. + if let Some(i_line) = import.lines().next() { + if let Some(crate_name) = i_line.split("::").next() { + if let Some(index) = result.iter().position(|s| s.starts_with(crate_name)) { + i = Some(index); + } + } + } + + match i { + Some(i) => { + let mut i_lines = import.lines(); + + while let Some(i_line) = i_lines.next() { + if (i_line.starts_with("use ") && i_line.ends_with(";")) + || (!i_line.starts_with("use ") && i_line != "};") + { + let mut r_lines = result[i].lines().peekable(); + + while let Some(r_line) = r_lines.next() { + if (r_line.starts_with("use ") && r_line.ends_with(";")) + || (!r_line.starts_with("use ") && r_line != "};") + { + if let Some(new_result) = + insert_import(i_line, r_line, r_lines.clone(), &result[i]) + { + result[i] = new_result; + break; + } + } + } + } + } + } + None => { + // Insert and sort alphabetically. + result.push(import.trim().to_string()); + result.sort(); + } + } + } + + result +} + +/// Returns a vector containing all the modules names in a line of an import. +fn modules_names(line: &str) -> Vec { + let mut chars = line.chars().peekable(); + let mut module_name = String::new(); + let mut names: Vec = Vec::new(); + + while let Some(c) = chars.next() { + module_name.push(c); + + if c == ':' { + if module_name.contains("::") { + if let Some(&next_char) = chars.peek() { + if next_char != '{' { + names.push(module_name.clone()); + module_name.clear(); + } + } + } + } else if module_name.contains(",") || module_name.contains("{") { + names.push(module_name.clone()); + module_name.clear(); + } else if c == ';' { + names.push(module_name.clone()); + module_name.clear(); + } + } + + names +} + +/// Removes the punctuation characters from the end of the name of a module. +fn rm_punct(name: &String) -> String { + let chars_to_remove = [';', ':', ',', '{', '}']; + + name.trim_end_matches(|c| chars_to_remove.contains(&c)) + .to_string() +} + +/// Checks which kind of insertion should be applied to the import, based on the +/// name of the module. +fn insert_import<'a, T>( + i_line: &str, + r_line: &str, + r_lines: std::iter::Peekable, + curr_result: &String, +) -> Option +where + T: Iterator, +{ + let i_line_names = modules_names(i_line); + let mut i_line_names_iter = i_line_names.iter().peekable(); + let mut i_name_punct = i_line_names_iter.next().unwrap().to_string(); + if i_name_punct.starts_with("use ") { + i_name_punct = i_line_names_iter.next().unwrap().to_string(); + } + let i_name = rm_punct(&i_name_punct); + let i_name = i_name.trim().to_string(); + + let r_line_names = modules_names(r_line); + let mut r_line_names_iter = r_line_names.iter().peekable(); + let mut r_name_punct = r_line_names_iter.next().unwrap().to_string(); + if r_name_punct.starts_with("use ") { + r_name_punct = r_line_names_iter.next().unwrap().to_string(); + } + let r_name = rm_punct(&r_name_punct); + let r_name = r_name.trim().to_string(); + + if i_name == r_name { + insert_on_line(i_line_names_iter, r_line_names_iter, r_line, curr_result) + } else if i_name < r_name { + insert_before_line(i_line, r_line, r_name_punct, curr_result) + } else { + insert_after_line(i_line, r_line, r_lines, curr_result) + } +} + +/// Inserts the import on the current line. +/// +/// For example, "use libafl::feedback::CrashFeedback;" and +/// "use libafl::feedback::TimeoutFeedback" will return: +/// "use libafl::feedback::{CrashFeedback, TimeoutFeedback};" +fn insert_on_line<'a, T>( + mut i_line_names_iter: std::iter::Peekable, + mut r_line_names_iter: std::iter::Peekable, + r_line: &str, + curr_result: &String, +) -> Option +where + T: Iterator + std::iter::DoubleEndedIterator, + T: Clone, +{ + let mut inserted = false; + let mut new_r_line = r_line.to_string(); + + while let Some(i_name_punct) = i_line_names_iter.peek() { + let i_name = rm_punct(&i_name_punct); + + if let Some(r_name_punct) = r_line_names_iter.peek() { + let r_name = rm_punct(&r_name_punct); + + if i_name < r_name { + if r_name_punct.ends_with(",") { + // Make it a multiple import and insert here. + let insert_string = rm_punct(&i_name_punct); + let (first, second) = + new_r_line.split_at(new_r_line.find(r_name_punct.as_str()).unwrap()); + let second = rm_punct(&second.to_string()); + let mut end = "},"; + + if new_r_line.ends_with(";") { + end = "};"; + } + new_r_line = format!("{}{}{}, {}{}", first, "{", insert_string, second, end); + inserted = true; + } else { + let mut r_line_names_iter_rev = r_line_names_iter.clone().rev(); + + if let Some(r_name_punct) = r_line_names_iter_rev.next() { + let insert_string = rm_punct(&i_name_punct); + let (first, second) = + new_r_line.split_at(new_r_line.find(r_name_punct).unwrap()); + let second = rm_punct(&second.to_string()); + let mut end = "},"; + + if new_r_line.ends_with(";") { + end = "};"; + } + new_r_line = + format!("{}{}{}, {}{}", first, "{", insert_string, second, end); + inserted = true; + } + } + } + } + // Advance the iterators + if let Some(r_name_punct) = r_line_names_iter.next() { + if let Some(i_name_punct) = i_line_names_iter.next() { + let r_name = rm_punct(&r_name_punct); + + if i_name > r_name { + // If there are no more elements, insert at the end. + if let None = r_line_names_iter.peek() { + let mut insert_string = i_name_punct.to_string(); + + while let Some(i_name_punct) = i_line_names_iter.next() { + insert_string.push_str(i_name_punct); + } + insert_string = rm_punct(&insert_string).trim_end().to_string(); + if r_name_punct.ends_with("},") || r_name_punct.ends_with("};") { + let (first, second) = + new_r_line.split_at(new_r_line.rfind("}").unwrap()); + + new_r_line = format!("{}, {}{}", first, insert_string, second); + } else { + let (first, second) = + r_line.split_at(r_line.find(r_name_punct.as_str()).unwrap()); + let second = second.trim_end().replace(";", ","); + let mut end = "},"; + + if new_r_line.ends_with(";") { + end = "};"; + } + new_r_line = + format!("{}{}{} {}{}", first, "{", second, insert_string, end); + } + let mut result_lines: Vec = + curr_result.lines().map(|line| line.to_string()).collect(); + + for (i, line) in curr_result.lines().enumerate() { + if line == r_line { + result_lines.remove(i); + result_lines.insert(i, new_r_line); + break; + } + } + + return Some(result_lines.join("\n")); + } + } + } + } + } + if inserted { + let mut result_lines: Vec = + curr_result.lines().map(|line| line.to_string()).collect(); + + for (i, line) in curr_result.lines().enumerate() { + if line == r_line { + result_lines.remove(i); + result_lines.insert(i, new_r_line); + break; + } + } + + return Some(result_lines.join("\n")); + } + + None +} + +/// Inserts the import before the current line. +/// +/// For example, "use libafl::executor::InProcessExecutor;" and +/// "use libafl::{ +/// feedback::TimeoutFeedback, +/// };", will return: +/// "use libafl::{ +/// executor::InProcessExecutor, +/// feedback::TimeoutFeedback, +/// };" +fn insert_before_line( + i_line: &str, + r_line: &str, + r_name_punct: String, + curr_result: &String, +) -> Option { + let new_result: String; + + if r_line.ends_with(";") { + // Change format to multiple import and insert as the first line. + let (first, second) = r_line.split_at(r_line.find("::").unwrap() + 2); + let insert_line = format_insert_line(i_line); + let mut first = first.to_string(); + let second = format_insert_line(second); + + first.push_str("{\n"); + new_result = format!("{}{}{}{}", first, insert_line, second, "};"); + } else { + // Simply insert the line here.. + let mut result_lines: Vec = + curr_result.lines().map(|line| line.to_string()).collect(); + + for (i, line) in curr_result.lines().enumerate() { + if line.contains(&r_name_punct) { + let insert_line = format_insert_line(i_line).replace("\n", ""); + + result_lines.insert(i, insert_line); + break; + } + } + new_result = result_lines.join("\n"); + } + + return Some(new_result); +} + +/// Inserts the import before the current line. +/// +/// For example, "use libafl::state::StdState;" and +/// "use libafl::{ +/// feedback::TimeoutFeedback, +/// };", will return: +/// "use libafl::{ +/// feedback::TimeoutFeedback, +/// state::StdState, +/// };" +fn insert_after_line<'a, T>( + i_line: &str, + r_line: &str, + mut r_lines: std::iter::Peekable, + curr_result: &String, +) -> Option +where + T: Iterator, +{ + if r_line.ends_with(";") { + // Change format to multiple import and insert as the last line. + let new_result: String; + let (first, second) = r_line.split_at(r_line.find("::").unwrap() + 2); + let insert_line = format_insert_line(i_line); + let mut first = first.to_string(); + let second = format_insert_line(second); + + first.push_str("{\n"); + new_result = format!("{}{}{}{}", first, second, insert_line, "};"); + + return Some(new_result); + } else { + if let Some(end) = r_lines.peek() { + if end == &"};" { + // If there are no more elements, insert as the last line. + let new_result: String; + let mut result_lines: Vec = + curr_result.lines().map(|line| line.to_string()).collect(); + let insert_line = format_insert_line(i_line).replace("\n", ""); + + result_lines.insert(result_lines.len() - 1, insert_line); + new_result = result_lines.join("\n"); + + return Some(new_result); + } + } + } + + None +} + +/// Formats the import line for insertion. +fn format_insert_line(line: &str) -> String { + let mut insert_line = line.trim().to_string(); + + if insert_line.starts_with("use ") { + let (_, second) = insert_line.split_once("::").unwrap(); + insert_line = second.to_string(); + } + insert_line.insert_str(0, &" ".repeat(4)); + insert_line.push('\n'); + insert_line = insert_line.replace(";", ","); + + insert_line +}