Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Commit

Permalink
Template creation command (#3531)
Browse files Browse the repository at this point in the history
* Tasks are selectable

* Almost there

* It works

* fmt

* remove dead code

* Remove unnecessary comments

* Improve instructions

* fix bug

* Add some dummy values for paths
  • Loading branch information
tevoinea authored Sep 28, 2023
1 parent e12b41e commit 552df45
Show file tree
Hide file tree
Showing 14 changed files with 470 additions and 20 deletions.
10 changes: 6 additions & 4 deletions src/agent/onefuzz-task/src/local/cmd.rs
Original file line number Diff line number Diff line change
@@ -1,26 +1,26 @@
// 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 {
#[cfg(any(target_os = "linux", target_os = "windows"))]
Coverage,
LibfuzzerFuzz,
Template,
CreateTemplate,
}

const TIMEOUT: &str = "timeout";
Expand All @@ -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();
Expand All @@ -62,6 +62,7 @@ pub async fn run(args: clap::ArgMatches) -> Result<()> {

template::launch(config, event_sender).await
}
Commands::CreateTemplate => create_template::run(),
}
});

Expand Down Expand Up @@ -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 {
Expand Down
15 changes: 14 additions & 1 deletion src/agent/onefuzz-task/src/local/coverage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,20 @@ pub struct Coverage {
}

#[async_trait]
impl Template for Coverage {
impl Template<Coverage> 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<Vec<SyncedDir>> = self
.readonly_inputs
Expand Down
285 changes: 285 additions & 0 deletions src/agent/onefuzz-task/src/local/create_template.rs
Original file line number Diff line number Diff line change
@@ -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<B: Backend>(terminal: &mut Terminal<B>, mut app: App) -> Result<Option<PathBuf>> {
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<ListElement>) -> Result<PathBuf> {
let tasks: Vec<TaskConfig> = 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<B: Backend>(f: &mut Frame<B>, 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<ListItem> = 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 <space> to select. Press <enter> 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<ListElement<'a>>,
}

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<ListElement> {
state: ListState,
items: Vec<ListElement>,
}

impl<T: Toggle> StatefulList<T> {
fn with_items(items: Vec<T>) -> StatefulList<T> {
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()
}
}
}
}
18 changes: 17 additions & 1 deletion src/agent/onefuzz-task/src/local/generic_analysis.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,23 @@ pub struct Analysis {
}

#[async_trait]
impl Template for Analysis {
impl Template<Analysis> 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?)
Expand Down
Loading

0 comments on commit 552df45

Please sign in to comment.