diff --git a/src/agent/onefuzz-task/src/local/cmd.rs b/src/agent/onefuzz-task/src/local/cmd.rs index eabefb71ee..cb800d445e 100644 --- a/src/agent/onefuzz-task/src/local/cmd.rs +++ b/src/agent/onefuzz-task/src/local/cmd.rs @@ -1,19 +1,18 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use super::{create_template, template}; #[cfg(any(target_os = "linux", target_os = "windows"))] use crate::local::coverage; use crate::local::{common::add_common_config, libfuzzer_fuzz, tui::TerminalUi}; use anyhow::{Context, Result}; + use clap::{Arg, ArgAction, Command}; use std::time::Duration; use std::{path::PathBuf, str::FromStr}; use strum::IntoEnumIterator; use strum_macros::{EnumIter, EnumString, IntoStaticStr}; use tokio::{select, time::timeout}; - -use super::template; - #[derive(Debug, PartialEq, Eq, EnumString, IntoStaticStr, EnumIter)] #[strum(serialize_all = "kebab-case")] enum Commands { @@ -21,6 +20,7 @@ enum Commands { Coverage, LibfuzzerFuzz, Template, + CreateTemplate, } const TIMEOUT: &str = "timeout"; @@ -43,7 +43,7 @@ pub async fn run(args: clap::ArgMatches) -> Result<()> { let sub_args = sub_args.clone(); - let terminal = if start_ui { + let terminal = if start_ui && command != Commands::CreateTemplate { Some(TerminalUi::init()?) } else { env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); @@ -62,6 +62,7 @@ pub async fn run(args: clap::ArgMatches) -> Result<()> { template::launch(config, event_sender).await } + Commands::CreateTemplate => create_template::run(), } }); @@ -116,6 +117,7 @@ pub fn args(name: &'static str) -> Command { .args(vec![Arg::new("config") .value_parser(value_parser!(std::path::PathBuf)) .required(true)]), + Commands::CreateTemplate => create_template::args(subcommand.into()), }; cmd = if add_common { diff --git a/src/agent/onefuzz-task/src/local/coverage.rs b/src/agent/onefuzz-task/src/local/coverage.rs index d091b70695..48e32cb861 100644 --- a/src/agent/onefuzz-task/src/local/coverage.rs +++ b/src/agent/onefuzz-task/src/local/coverage.rs @@ -148,7 +148,20 @@ pub struct Coverage { } #[async_trait] -impl Template for Coverage { +impl Template for Coverage { + fn example_values() -> Coverage { + Coverage { + target_exe: PathBuf::from("path_to_your_exe"), + target_env: HashMap::new(), + target_options: vec![], + target_timeout: None, + module_allowlist: None, + source_allowlist: None, + input_queue: Some(PathBuf::from("path_to_your_inputs")), + readonly_inputs: vec![PathBuf::from("path_to_readonly_inputs")], + coverage: PathBuf::from("path_to_where_you_want_coverage_to_be_output"), + } + } async fn run(&self, context: &RunContext) -> Result<()> { let ri: Result> = self .readonly_inputs diff --git a/src/agent/onefuzz-task/src/local/create_template.rs b/src/agent/onefuzz-task/src/local/create_template.rs new file mode 100644 index 0000000000..474b677ad0 --- /dev/null +++ b/src/agent/onefuzz-task/src/local/create_template.rs @@ -0,0 +1,285 @@ +use crate::local::template::CommonProperties; + +use super::template::{TaskConfig, TaskConfigDiscriminants, TaskGroup}; +use anyhow::Result; +use clap::Command; +use std::str::FromStr; +use std::{ + io, + path::{Path, PathBuf}, +}; + +use strum::VariantNames; + +use crate::local::{ + coverage::Coverage, generic_analysis::Analysis, generic_crash_report::CrashReport, + generic_generator::Generator, libfuzzer::LibFuzzer, + libfuzzer_crash_report::LibfuzzerCrashReport, libfuzzer_merge::LibfuzzerMerge, + libfuzzer_regression::LibfuzzerRegression, libfuzzer_test_input::LibfuzzerTestInput, + template::Template, test_input::TestInput, +}; + +use crossterm::{ + event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind}, + execute, + terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, +}; +use tui::{prelude::*, widgets::*}; + +pub fn args(name: &'static str) -> Command { + Command::new(name).about("interactively create a template") +} + +pub fn run() -> Result<()> { + // setup terminal + enable_raw_mode()?; + let mut stdout = io::stdout(); + execute!(stdout, EnterAlternateScreen, EnableMouseCapture)?; + let backend = CrosstermBackend::new(stdout); + let mut terminal = Terminal::new(backend)?; + + // create app and run it + let app = App::new(); + let res = run_app(&mut terminal, app); + + // restore terminal + disable_raw_mode()?; + execute!( + terminal.backend_mut(), + LeaveAlternateScreen, + DisableMouseCapture + )?; + terminal.show_cursor()?; + + match res { + Ok(None) => { /* user quit, do nothing */ } + Ok(Some(path)) => match path.canonicalize() { + Ok(canonical_path) => println!("Wrote the template to: {:?}", canonical_path), + _ => println!("Wrote the template to: {:?}", path), + }, + Err(e) => println!("Failed to write template due to {}", e), + } + + Ok(()) +} + +fn run_app(terminal: &mut Terminal, mut app: App) -> Result> { + loop { + terminal.draw(|f| ui(f, &mut app))?; + if let Event::Key(key) = event::read()? { + if key.kind == KeyEventKind::Press { + match key.code { + KeyCode::Char('q') => return Ok(None), + KeyCode::Char(' ') => app.items.toggle(), + KeyCode::Down => app.items.next(), + KeyCode::Up => app.items.previous(), + KeyCode::Enter => { + return match generate_template(app.items.items) { + Ok(p) => Ok(Some(p)), + Err(e) => Err(e), + } + } + _ => {} + } + } + } + } +} + +fn generate_template(items: Vec) -> Result { + let tasks: Vec = items + .iter() + .filter(|item| item.is_included) + .filter_map(|list_element| { + match TaskConfigDiscriminants::from_str(list_element.task_type) { + Err(e) => { + error!( + "Failed to match task config {:?} - {}", + list_element.task_type, e + ); + None + } + Ok(t) => match t { + TaskConfigDiscriminants::LibFuzzer => { + Some(TaskConfig::LibFuzzer(LibFuzzer::example_values())) + } + TaskConfigDiscriminants::Analysis => { + Some(TaskConfig::Analysis(Analysis::example_values())) + } + TaskConfigDiscriminants::Coverage => { + Some(TaskConfig::Coverage(Coverage::example_values())) + } + TaskConfigDiscriminants::CrashReport => { + Some(TaskConfig::CrashReport(CrashReport::example_values())) + } + TaskConfigDiscriminants::Generator => { + Some(TaskConfig::Generator(Generator::example_values())) + } + TaskConfigDiscriminants::LibfuzzerCrashReport => Some( + TaskConfig::LibfuzzerCrashReport(LibfuzzerCrashReport::example_values()), + ), + TaskConfigDiscriminants::LibfuzzerMerge => { + Some(TaskConfig::LibfuzzerMerge(LibfuzzerMerge::example_values())) + } + TaskConfigDiscriminants::LibfuzzerRegression => Some( + TaskConfig::LibfuzzerRegression(LibfuzzerRegression::example_values()), + ), + TaskConfigDiscriminants::LibfuzzerTestInput => Some( + TaskConfig::LibfuzzerTestInput(LibfuzzerTestInput::example_values()), + ), + TaskConfigDiscriminants::TestInput => { + Some(TaskConfig::TestInput(TestInput::example_values())) + } + TaskConfigDiscriminants::Radamsa => Some(TaskConfig::Radamsa), + }, + } + }) + .collect(); + + let definition = TaskGroup { + common: CommonProperties { + setup_dir: None, + extra_setup_dir: None, + extra_dir: None, + create_job_dir: false, + }, + tasks, + }; + + let filename = "template"; + let mut filepath = format!("./{}.yaml", filename); + let mut output_file = Path::new(&filepath); + let mut counter = 0; + while output_file.exists() { + filepath = format!("./{}-{}.yaml", filename, counter); + output_file = Path::new(&filepath); + counter += 1; + } + + std::fs::write(output_file, serde_yaml::to_string(&definition)?)?; + + Ok(output_file.into()) +} + +fn ui(f: &mut Frame, app: &mut App) { + let areas = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Percentage(100)]) + .split(f.size()); + // Iterate through all elements in the `items` app and append some debug text to it. + let items: Vec = app + .items + .items + .iter() + .map(|list_element| { + let title = if list_element.is_included { + format!("✅ {}", list_element.task_type) + } else { + list_element.task_type.to_string() + }; + ListItem::new(title).style(Style::default().fg(Color::Black).bg(Color::White)) + }) + .collect(); + + // Create a List from all list items and highlight the currently selected one + let items = List::new(items) + .block( + Block::default() + .borders(Borders::ALL) + .title("Select which tasks you want to include in the template. Use ⬆/⬇ to navigate and to select. Press when you're done."), + ) + .highlight_style( + Style::default() + .bg(Color::LightGreen) + .add_modifier(Modifier::BOLD), + ) + .highlight_symbol(">> "); + + // We can now render the item list + f.render_stateful_widget(items, areas[0], &mut app.items.state); +} + +struct ListElement<'a> { + pub task_type: &'a str, + pub is_included: bool, +} + +pub trait Toggle { + fn toggle(&mut self) {} +} + +impl<'a> Toggle for ListElement<'a> { + fn toggle(&mut self) { + self.is_included = !self.is_included + } +} + +struct App<'a> { + items: StatefulList>, +} + +impl<'a> App<'a> { + fn new() -> App<'a> { + App { + items: StatefulList::with_items( + TaskConfig::VARIANTS + .iter() + .map(|name| ListElement { + task_type: name, + is_included: false, + }) + .collect(), + ), + } + } +} + +struct StatefulList { + state: ListState, + items: Vec, +} + +impl StatefulList { + fn with_items(items: Vec) -> StatefulList { + StatefulList { + state: ListState::default(), + items, + } + } + + fn next(&mut self) { + let i = match self.state.selected() { + Some(i) => { + if self.items.first().is_some() { + (i + 1) % self.items.len() + } else { + 0 + } + } + None => 0, + }; + self.state.select(Some(i)); + } + + fn previous(&mut self) { + let i = match self.state.selected() { + Some(i) => { + if i == 0 { + self.items.len() - 1 + } else { + i - 1 + } + } + None => 0, + }; + self.state.select(Some(i)); + } + + fn toggle(&mut self) { + if let Some(index) = self.state.selected() { + if let Some(element) = self.items.get_mut(index) { + element.toggle() + } + } + } +} diff --git a/src/agent/onefuzz-task/src/local/generic_analysis.rs b/src/agent/onefuzz-task/src/local/generic_analysis.rs index 429e7b0e3b..cbb31a1ff9 100644 --- a/src/agent/onefuzz-task/src/local/generic_analysis.rs +++ b/src/agent/onefuzz-task/src/local/generic_analysis.rs @@ -27,7 +27,23 @@ pub struct Analysis { } #[async_trait] -impl Template for Analysis { +impl Template for Analysis { + fn example_values() -> Analysis { + Analysis { + analyzer_exe: String::new(), + analyzer_options: vec![], + analyzer_env: HashMap::new(), + target_exe: PathBuf::from("path_to_your_exe"), + target_options: vec![], + input_queue: Some(PathBuf::from("path_to_your_inputs")), + crashes: Some(PathBuf::from("path_where_crashes_written")), + analysis: PathBuf::new(), + tools: None, + reports: Some(PathBuf::from("path_where_reports_written")), + unique_reports: Some(PathBuf::from("path_where_reports_written")), + no_repro: Some(PathBuf::from("path_where_no_repro_reports_written")), + } + } async fn run(&self, context: &RunContext) -> Result<()> { let input_q = if let Some(w) = &self.input_queue { Some(context.monitor_dir(w).await?) diff --git a/src/agent/onefuzz-task/src/local/generic_crash_report.rs b/src/agent/onefuzz-task/src/local/generic_crash_report.rs index 347a8cac76..91dec1ae44 100644 --- a/src/agent/onefuzz-task/src/local/generic_crash_report.rs +++ b/src/agent/onefuzz-task/src/local/generic_crash_report.rs @@ -39,7 +39,25 @@ pub struct CrashReport { minimized_stack_depth: Option, } #[async_trait] -impl Template for CrashReport { +impl Template for CrashReport { + fn example_values() -> CrashReport { + CrashReport { + target_exe: PathBuf::from("path_to_your_exe"), + target_options: vec![], + target_env: HashMap::new(), + input_queue: Some(PathBuf::from("path_to_your_inputs")), + crashes: Some(PathBuf::from("path_where_crashes_written")), + reports: Some(PathBuf::from("path_where_reports_written")), + unique_reports: Some(PathBuf::from("path_where_reports_written")), + no_repro: Some(PathBuf::from("path_where_no_repro_reports_written")), + target_timeout: None, + check_asan_log: true, + check_debugger: true, + check_retry_count: 5, + check_queue: false, + minimized_stack_depth: None, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let input_q_fut: OptionFuture<_> = self .input_queue diff --git a/src/agent/onefuzz-task/src/local/generic_generator.rs b/src/agent/onefuzz-task/src/local/generic_generator.rs index ae9f6a3cc6..3c26af4cf8 100644 --- a/src/agent/onefuzz-task/src/local/generic_generator.rs +++ b/src/agent/onefuzz-task/src/local/generic_generator.rs @@ -35,7 +35,26 @@ pub struct Generator { } #[async_trait] -impl Template for Generator { +impl Template for Generator { + fn example_values() -> Generator { + Generator { + generator_exe: String::new(), + generator_env: HashMap::new(), + generator_options: vec![], + readonly_inputs: vec![PathBuf::from("path_to_readonly_inputs")], + crashes: PathBuf::new(), + tools: None, + target_exe: PathBuf::from("path_to_your_exe"), + target_env: HashMap::new(), + target_options: vec![], + target_timeout: None, + check_asan_log: true, + check_debugger: true, + check_retry_count: 5, + rename_output: false, + ensemble_sync_delay: None, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let generator_config = crate::tasks::fuzz::generator::Config { generator_exe: self.generator_exe.clone(), diff --git a/src/agent/onefuzz-task/src/local/libfuzzer.rs b/src/agent/onefuzz-task/src/local/libfuzzer.rs index 433636be1c..472a6ae9e8 100644 --- a/src/agent/onefuzz-task/src/local/libfuzzer.rs +++ b/src/agent/onefuzz-task/src/local/libfuzzer.rs @@ -32,7 +32,22 @@ pub struct LibFuzzer { } #[async_trait] -impl Template for LibFuzzer { +impl Template for LibFuzzer { + fn example_values() -> LibFuzzer { + LibFuzzer { + inputs: PathBuf::new(), + readonly_inputs: vec![PathBuf::from("path_to_readonly_inputs")], + crashes: PathBuf::new(), + crashdumps: None, + target_exe: PathBuf::from("path_to_your_exe"), + target_env: HashMap::new(), + target_options: vec![], + target_workers: None, + ensemble_sync_delay: None, + check_fuzzer_help: true, + expect_crash_on_failure: true, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let ri: Result> = self .readonly_inputs diff --git a/src/agent/onefuzz-task/src/local/libfuzzer_crash_report.rs b/src/agent/onefuzz-task/src/local/libfuzzer_crash_report.rs index 04ba4f9225..9de1fc66ce 100644 --- a/src/agent/onefuzz-task/src/local/libfuzzer_crash_report.rs +++ b/src/agent/onefuzz-task/src/local/libfuzzer_crash_report.rs @@ -36,7 +36,24 @@ pub struct LibfuzzerCrashReport { } #[async_trait] -impl Template for LibfuzzerCrashReport { +impl Template for LibfuzzerCrashReport { + fn example_values() -> LibfuzzerCrashReport { + LibfuzzerCrashReport { + target_exe: PathBuf::from("path_to_your_exe"), + target_env: HashMap::new(), + target_options: vec![], + target_timeout: None, + input_queue: Some(PathBuf::from("path_to_your_inputs")), + crashes: Some(PathBuf::from("path_where_crashes_written")), + reports: Some(PathBuf::from("path_where_reports_written")), + unique_reports: Some(PathBuf::from("path_where_reports_written")), + no_repro: Some(PathBuf::from("path_where_no_repro_reports_written")), + check_fuzzer_help: true, + check_retry_count: 5, + minimized_stack_depth: None, + check_queue: true, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let input_q_fut: OptionFuture<_> = self .input_queue diff --git a/src/agent/onefuzz-task/src/local/libfuzzer_merge.rs b/src/agent/onefuzz-task/src/local/libfuzzer_merge.rs index 4b3e4ce58f..d4915e6b4c 100644 --- a/src/agent/onefuzz-task/src/local/libfuzzer_merge.rs +++ b/src/agent/onefuzz-task/src/local/libfuzzer_merge.rs @@ -27,7 +27,19 @@ pub struct LibfuzzerMerge { } #[async_trait] -impl Template for LibfuzzerMerge { +impl Template for LibfuzzerMerge { + fn example_values() -> LibfuzzerMerge { + LibfuzzerMerge { + target_exe: PathBuf::from("path_to_your_exe"), + target_env: HashMap::new(), + target_options: vec![], + input_queue: Some(PathBuf::from("path_to_your_inputs")), + inputs: vec![], + unique_inputs: PathBuf::new(), + preserve_existing_outputs: true, + check_fuzzer_help: true, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let input_q_fut: OptionFuture<_> = self .input_queue diff --git a/src/agent/onefuzz-task/src/local/libfuzzer_regression.rs b/src/agent/onefuzz-task/src/local/libfuzzer_regression.rs index 3fbb9f0bd6..b53fb84c22 100644 --- a/src/agent/onefuzz-task/src/local/libfuzzer_regression.rs +++ b/src/agent/onefuzz-task/src/local/libfuzzer_regression.rs @@ -40,7 +40,25 @@ pub struct LibfuzzerRegression { } #[async_trait] -impl Template for LibfuzzerRegression { +impl Template for LibfuzzerRegression { + fn example_values() -> LibfuzzerRegression { + LibfuzzerRegression { + target_exe: PathBuf::from("path_to_your_exe"), + target_options: vec![], + target_env: HashMap::new(), + target_timeout: None, + crashes: PathBuf::new(), + regression_reports: PathBuf::new(), + report_list: None, + unique_reports: Some(PathBuf::from("path_where_reports_written")), + reports: Some(PathBuf::from("path_where_reports_written")), + no_repro: Some(PathBuf::from("path_where_no_repro_reports_written")), + readonly_inputs: None, + check_fuzzer_help: true, + check_retry_count: 5, + minimized_stack_depth: None, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let libfuzzer_regression = crate::tasks::regression::libfuzzer::Config { target_exe: self.target_exe.clone(), diff --git a/src/agent/onefuzz-task/src/local/libfuzzer_test_input.rs b/src/agent/onefuzz-task/src/local/libfuzzer_test_input.rs index 5bef2347f7..88c3cd1a3d 100644 --- a/src/agent/onefuzz-task/src/local/libfuzzer_test_input.rs +++ b/src/agent/onefuzz-task/src/local/libfuzzer_test_input.rs @@ -24,7 +24,21 @@ pub struct LibfuzzerTestInput { } #[async_trait] -impl Template for LibfuzzerTestInput { +impl Template for LibfuzzerTestInput { + fn example_values() -> LibfuzzerTestInput { + LibfuzzerTestInput { + input: PathBuf::new(), + target_exe: PathBuf::from("path_to_your_exe"), + target_options: vec![], + target_env: HashMap::new(), + setup_dir: PathBuf::new(), + extra_setup_dir: None, + extra_output_dir: None, + target_timeout: None, + check_retry_count: 5, + minimized_stack_depth: None, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let c = self.clone(); let t = tokio::spawn(async move { diff --git a/src/agent/onefuzz-task/src/local/mod.rs b/src/agent/onefuzz-task/src/local/mod.rs index 385ff8ffcd..6020cb0fa6 100644 --- a/src/agent/onefuzz-task/src/local/mod.rs +++ b/src/agent/onefuzz-task/src/local/mod.rs @@ -5,6 +5,7 @@ pub mod cmd; pub mod common; #[cfg(any(target_os = "linux", target_os = "windows"))] pub mod coverage; +pub mod create_template; pub mod generic_analysis; pub mod generic_crash_report; pub mod generic_generator; diff --git a/src/agent/onefuzz-task/src/local/template.rs b/src/agent/onefuzz-task/src/local/template.rs index 64b342744d..3393edd89a 100644 --- a/src/agent/onefuzz-task/src/local/template.rs +++ b/src/agent/onefuzz-task/src/local/template.rs @@ -5,6 +5,7 @@ use path_absolutize::Absolutize; use serde::Deserialize; use std::path::{Path, PathBuf}; use storage_queue::QueueClient; +use strum_macros::{EnumDiscriminants, EnumString, EnumVariantNames}; use tokio::{sync::Mutex, task::JoinHandle}; use url::Url; use uuid::Uuid; @@ -27,14 +28,14 @@ use schemars::JsonSchema; #[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] pub struct TaskGroup { #[serde(flatten)] - common: CommonProperties, + pub common: CommonProperties, /// The list of tasks - tasks: Vec, + pub tasks: Vec, } #[derive(Debug, Deserialize, Serialize, Clone, JsonSchema)] -struct CommonProperties { +pub struct CommonProperties { pub setup_dir: Option, pub extra_setup_dir: Option, pub extra_dir: Option, @@ -42,9 +43,10 @@ struct CommonProperties { pub create_job_dir: bool, } -#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema)] +#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, EnumVariantNames, EnumDiscriminants)] +#[strum_discriminants(derive(EnumString))] #[serde(tag = "type")] -enum TaskConfig { +pub enum TaskConfig { LibFuzzer(LibFuzzer), Analysis(Analysis), Coverage(Coverage), @@ -61,7 +63,8 @@ enum TaskConfig { } #[async_trait] -pub trait Template { +pub trait Template { + fn example_values() -> T; async fn run(&self, context: &RunContext) -> Result<()>; } diff --git a/src/agent/onefuzz-task/src/local/test_input.rs b/src/agent/onefuzz-task/src/local/test_input.rs index b8027a7f41..0018494ec0 100644 --- a/src/agent/onefuzz-task/src/local/test_input.rs +++ b/src/agent/onefuzz-task/src/local/test_input.rs @@ -28,7 +28,24 @@ pub struct TestInput { } #[async_trait] -impl Template for TestInput { +impl Template for TestInput { + fn example_values() -> TestInput { + TestInput { + input: PathBuf::new(), + target_exe: PathBuf::from("path_to_your_exe"), + target_options: vec![], + target_env: HashMap::new(), + setup_dir: PathBuf::new(), + extra_setup_dir: None, + task_id: Uuid::new_v4(), + job_id: Uuid::new_v4(), + target_timeout: None, + check_retry_count: 5, + check_asan_log: true, + check_debugger: true, + minimized_stack_depth: None, + } + } async fn run(&self, context: &RunContext) -> Result<()> { let c = self.clone(); let t = tokio::spawn(async move {