From 8950f3f356318c8be9e8b68474f248d487df2bff Mon Sep 17 00:00:00 2001 From: Jon-Becker Date: Wed, 22 Nov 2023 19:27:07 -0600 Subject: [PATCH] fix(cli): standardize `--output` flag paths, opts --- cli/src/main.rs | 163 ++++++++++++---------------------- cli/src/output.rs | 76 ++++++++++++++++ core/src/cfg/mod.rs | 8 +- core/src/cfg/output.rs | 73 +-------------- core/src/decompile/mod.rs | 2 +- core/src/disassemble/mod.rs | 2 +- core/src/dump/mod.rs | 2 +- core/src/snapshot/mod.rs | 4 +- core/src/snapshot/util/csv.rs | 9 +- core/tests/test_cfg.rs | 4 - 10 files changed, 147 insertions(+), 196 deletions(-) create mode 100644 cli/src/output.rs diff --git a/cli/src/main.rs b/cli/src/main.rs index c738fbab..53abc911 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,5 +1,8 @@ +pub(crate) mod output; + use backtrace::Backtrace; -use std::{env, io, panic}; +use output::build_output_path; +use std::{io, panic}; use clap::{Parser, Subcommand}; use colored::Colorize; @@ -10,25 +13,21 @@ use crossterm::{ }; use heimdall_cache::{cache, CacheArgs}; -use heimdall_common::{ - constants::ADDRESS_REGEX, - ether::rpc, - utils::{ - io::{ - file::{write_file, write_lines_to_file}, - logging::Logger, - }, - version::{current_version, remote_version}, +use heimdall_common::utils::{ + io::{ + file::{write_file, write_lines_to_file}, + logging::Logger, }, + version::{current_version, remote_version}, }; use heimdall_config::{config, get_config, ConfigArgs}; use heimdall_core::{ - cfg::{cfg, output::write_cfg_to_file, CFGArgs}, + cfg::{cfg, output::build_cfg, CFGArgs}, decode::{decode, DecodeArgs}, decompile::{decompile, out::abi::ABIStructure, DecompilerArgs}, disassemble::{disassemble, DisassemblerArgs}, dump::{dump, DumpArgs}, - snapshot::{snapshot, util::csv::generate_and_write_contract_csv, SnapshotArgs}, + snapshot::{snapshot, util::csv::generate_csv, SnapshotArgs}, }; use tui::{backend::CrosstermBackend, Terminal}; @@ -68,7 +67,7 @@ pub enum Subcommands { Dump(DumpArgs), #[clap( name = "snapshot", - about = "Infer function information from bytecode, including access control, gas + about = "Infer functiogn information from bytecode, including access control, gas consumption, storage accesses, event emissions, and more" )] Snapshot(SnapshotArgs), @@ -99,15 +98,8 @@ async fn main() -> Result<(), Box> { })); let configuration = get_config(); - - // get the current working directory - let mut output_path = env::current_dir()?.into_os_string().into_string().unwrap(); - match args.sub { Subcommands::Disassemble(mut cmd) => { - // get specified output path - output_path.push_str(&format!("/{}", cmd.output)); - // if the user has not specified a rpc url, use the default if cmd.rpc_url.as_str() == "" { cmd.rpc_url = configuration.rpc_url; @@ -116,17 +108,14 @@ async fn main() -> Result<(), Box> { let assembly = disassemble(cmd.clone()).await?; if cmd.output == "print" { + // TODO: use `less` println!("{}", assembly); } else { - let (dir_path, filename) = if ADDRESS_REGEX.is_match(&cmd.target).unwrap() { - (format!("{}/{}", output_path, &cmd.target), "disassembled.asm") - } else { - (format!("{}/local", output_path), "disassembled.asm") - }; - - std::fs::create_dir_all(&dir_path).expect("Failed to create output directory"); - let full_path = format!("{}/{}", dir_path, filename); - write_file(&full_path, &assembly); + let output_path = + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "disassembled.asm") + .await?; + + write_file(&output_path, &assembly); } } @@ -138,43 +127,22 @@ async fn main() -> Result<(), Box> { let result = decompile(cmd.clone()).await?; - // get specified output path - output_path.push_str(&format!("/{}", cmd.output)); - if cmd.output == "print" { if let Some(abi) = &result.abi { - println!("ABI: {}", serde_json::to_string_pretty(abi).unwrap()); + println!("ABI:\n\n{}\n", serde_json::to_string_pretty(abi).unwrap()); } if let Some(source) = &result.source { - println!("Source: {}", source); + println!("Source:\n\n{}\n", source); } } else { - // write to file - let (dir_path, abi_filename, solidity_filename, yul_filename) = - if ADDRESS_REGEX.is_match(&cmd.target).unwrap() { - let chain_id = rpc::chain_id(&cmd.rpc_url).await.unwrap(); - ( - format!("{}/{}/{}", output_path, chain_id, cmd.target), - "abi.json", - "decompiled.sol", - "decompiled.yul", - ) - } else { - ( - format!("{}/local", output_path), - "abi.json", - "decompiled.sol", - "decompiled.yul", - ) - }; - - std::fs::create_dir_all(&dir_path).expect("Failed to create output directory"); - + // write the contract ABI if let Some(abi) = result.abi { - // write the ABI to a file - let full_path = format!("{}/{}", dir_path, abi_filename); + let output_path = + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "abi.json") + .await?; + write_file( - &full_path, + &output_path, &format!( "[{}]", abi.iter() @@ -196,13 +164,17 @@ async fn main() -> Result<(), Box> { ), ); } + + // write the contract source if let Some(source) = &result.source { - let full_path = if cmd.include_solidity { - format!("{}/{}", dir_path, solidity_filename) + let output_path = if cmd.include_solidity { + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "decompiled.sol") + .await? } else { - format!("{}/{}", dir_path, yul_filename) + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "decompiled.yul") + .await? }; - write_file(&full_path, source); + write_file(&output_path, source); } } } @@ -231,20 +203,15 @@ async fn main() -> Result<(), Box> { } let cfg = cfg(cmd.clone()).await?; + let stringified_dot = build_cfg(&cfg, &cmd); - // get specified output path - output_path.push_str(&format!("/{}", cmd.output)); - - // write to file - let dir_path = if ADDRESS_REGEX.is_match(&cmd.target).unwrap() { - let chain_id = rpc::chain_id(&cmd.rpc_url).await.unwrap(); - format!("{}/{}/{}", output_path, chain_id, cmd.target) + if cmd.output == "print" { + println!("{}", stringified_dot); } else { - format!("{}/local", output_path) - }; - - std::fs::create_dir_all(&dir_path).expect("Failed to create output directory"); - write_cfg_to_file(&cfg, &cmd, dir_path); + let output_path = + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "cfg.dot").await?; + write_file(&output_path, &stringified_dot); + } } Subcommands::Dump(mut cmd) => { @@ -261,9 +228,6 @@ async fn main() -> Result<(), Box> { let result = dump(cmd.clone()).await?; let mut lines = Vec::new(); - // get specified output path - output_path.push_str(&format!("/{}", cmd.output)); - // add header lines.push(String::from("last_modified,alias,slot,decoded_type,value")); @@ -280,16 +244,9 @@ async fn main() -> Result<(), Box> { println!("{}", line); } } else { - // write to file - if ADDRESS_REGEX.is_match(&cmd.target).unwrap() { - output_path.push_str(&format!( - "/{}/{}/dump.csv", - rpc::chain_id(&cmd.rpc_url).await.unwrap(), - &cmd.target - )); - } else { - output_path.push_str("/local/dump.csv"); - } + let output_path = + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "dump.csv").await?; + write_lines_to_file(&output_path, lines); } } @@ -301,27 +258,25 @@ async fn main() -> Result<(), Box> { } let snapshot_result = snapshot(cmd.clone()).await?; - - // get specified output path - output_path.push_str(&format!("/{}", cmd.output)); - - // write to file - let dir_path = if ADDRESS_REGEX.is_match(&cmd.target).unwrap() { - let chain_id = rpc::chain_id(&cmd.rpc_url).await.unwrap(); - format!("{}/{}/{}", output_path, chain_id, cmd.target) - } else { - format!("{}/local", output_path) - }; - - std::fs::create_dir_all(&dir_path).expect("Failed to create output directory"); - let full_path = format!("{}/snapshot.csv", dir_path); - generate_and_write_contract_csv( + let csv_lines = generate_csv( &snapshot_result.snapshots, &snapshot_result.resolved_errors, &snapshot_result.resolved_events, - &full_path, - ) + ); + + if cmd.output == "print" { + for line in &csv_lines { + println!("{}", line); + } + } else { + let output_path = + build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "snapshot.csv") + .await?; + + write_lines_to_file(&output_path, csv_lines); + } } + Subcommands::Config(cmd) => { config(cmd); } diff --git a/cli/src/output.rs b/cli/src/output.rs new file mode 100644 index 00000000..e9f37823 --- /dev/null +++ b/cli/src/output.rs @@ -0,0 +1,76 @@ +use std::env; + +use heimdall_common::{constants::ADDRESS_REGEX, ether::rpc}; + +/// build a standardized output path for the given parameters. follows the following cases: +/// - if `output` is `print`, return `None` +/// - if `output` is the default value (`output`) +/// - if `target` is a contract_address, return `/output/{chain_id}/{target}/{filename}` +/// - if `target` is a file or raw bytes, return `/output/local/{filename}` +/// - if `output` is specified, return `/{output}/{filename}` +pub async fn build_output_path( + output: &str, + target: &str, + rpc_url: &str, + filename: &str, +) -> Result> { + // if output is the default value, build a path based on the target + if output == "output" { + // get the current working directory + let cwd = env::current_dir()?.into_os_string().into_string().unwrap(); + + if ADDRESS_REGEX.is_match(target)? { + let chain_id = rpc::chain_id(rpc_url).await?; + return Ok(format!("{}/output/{}/{}/{}", cwd, chain_id, target, filename)) + } else { + return Ok(format!("{}/output/local/{}", cwd, filename)) + } + } + + // output is specified, return the path + Ok(format!("{}/{}", output, filename)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_output_default_address() { + let output = "output"; + let target = "0x0000000000000000000000000000000000000001"; + let rpc_url = "https://eth.llamarpc.com"; + let filename = "cfg.dot"; + + let path = build_output_path(output, target, rpc_url, filename).await; + assert!(path.is_ok()); + assert!(path + .unwrap() + .ends_with("/output/1/0x0000000000000000000000000000000000000001/cfg.dot")); + } + + #[tokio::test] + async fn test_output_default_local() { + let output = "output"; + let target = + "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000"; + let rpc_url = "https://eth.llamarpc.com"; + let filename = "cfg.dot"; + + let path = build_output_path(output, target, rpc_url, filename).await; + assert!(path.is_ok()); + assert!(path.unwrap().ends_with("/output/local/cfg.dot")); + } + + #[tokio::test] + async fn test_output_specified() { + let output = "/some_dir"; + let target = "0x0000000000000000000000000000000000000001"; + let rpc_url = "https://eth.llamarpc.com"; + let filename = "cfg.dot"; + + let path = build_output_path(output, target, rpc_url, filename).await; + assert!(path.is_ok()); + assert_eq!(path.unwrap(), "/some_dir/cfg.dot".to_string()); + } +} diff --git a/core/src/cfg/mod.rs b/core/src/cfg/mod.rs index e4d4afe6..bcb7af1a 100644 --- a/core/src/cfg/mod.rs +++ b/core/src/cfg/mod.rs @@ -44,18 +44,13 @@ pub struct CFGArgs { #[clap(long, short)] pub default: bool, - /// Specify a format (other than dot) to output the CFG in. - /// For example, `--format svg` will output a SVG image of the CFG. - #[clap(long = "format", short, default_value = "", hide_default_value = true)] - pub format: String, - /// Color the edges of the graph based on the JUMPI condition. /// This is useful for visualizing the flow of if statements. #[clap(long = "color-edges", short)] pub color_edges: bool, /// The output directory to write the output to or 'print' to print to the console - #[clap(long = "output", short = 'o', default_value = "")] + #[clap(long = "output", short = 'o', default_value = "output", hide_default_value = true)] pub output: String, } @@ -66,7 +61,6 @@ impl CFGArgsBuilder { verbose: Some(clap_verbosity_flag::Verbosity::new(0, 1)), rpc_url: Some(String::new()), default: Some(true), - format: Some(String::new()), color_edges: Some(false), output: Some(String::new()), } diff --git a/core/src/cfg/output.rs b/core/src/cfg/output.rs index 1219b834..ce9a9509 100644 --- a/core/src/cfg/output.rs +++ b/core/src/cfg/output.rs @@ -1,23 +1,9 @@ -use std::{process::Command, time::Duration}; - -use heimdall_common::utils::io::{file::write_file, logging::Logger}; -use indicatif::ProgressBar; use petgraph::{dot::Dot, graph::Graph}; use super::CFGArgs; /// Write the generated CFG to a file in the `dot` graphviz format. -pub fn write_cfg_to_file(contract_cfg: &Graph, args: &CFGArgs, output_dir: String) { - // get a new logger - let logger = Logger::default(); - - // get a new progress bar - let progress_bar = ProgressBar::new_spinner(); - progress_bar.enable_steady_tick(Duration::from_millis(100)); - progress_bar.set_style(logger.info_spinner()); - progress_bar.set_message("writing CFG .dot file".to_string()); - - let dot_output_path = format!("{output_dir}/cfg.dot"); +pub fn build_cfg(contract_cfg: &Graph, args: &CFGArgs) -> String { let output = format!("{}", Dot::with_config(&contract_cfg, &[])); // find regex matches and replace @@ -38,60 +24,5 @@ pub fn write_cfg_to_file(contract_cfg: &Graph, args: &CFGArgs, o output = output.replace("[ label = \"\" ]", "[]"); - write_file(&dot_output_path, &output); - - progress_bar.suspend(|| { - logger.success(&format!("wrote generated dot to '{}' .", &dot_output_path)); - }); - - if !args.format.is_empty() { - // check for graphviz - match Command::new("dot").spawn() { - Ok(_) => { - progress_bar.set_message(format!("generating CFG .{} file", &args.format)); - - let image_output_path = format!("{}/cfg.{}", output_dir, &args.format); - match Command::new("dot").arg("-T").arg(&args.format).arg(&dot_output_path).output() - { - Ok(output) => { - match String::from_utf8(output.stdout) { - Ok(output) => { - // write the output - write_file(&image_output_path, &output); - progress_bar.suspend(|| { - logger.success(&format!( - "wrote generated {} to '{}' .", - &args.format, &image_output_path - )); - }); - } - Err(_) => { - progress_bar.suspend(|| { - logger.error(&format!( - "graphviz failed to generate {} file.", - &args.format - )); - }); - } - } - } - Err(_) => { - progress_bar.suspend(|| { - logger.error(&format!( - "graphviz failed to generate {} file.", - &args.format - )); - }); - } - } - } - Err(_) => { - progress_bar.suspend(|| { - logger.error("graphviz doesn't appear to be installed. please install graphviz to generate images."); - }); - } - } - } - - progress_bar.finish_and_clear(); + output } diff --git a/core/src/decompile/mod.rs b/core/src/decompile/mod.rs index 12966d18..fa005b99 100644 --- a/core/src/decompile/mod.rs +++ b/core/src/decompile/mod.rs @@ -73,7 +73,7 @@ pub struct DecompilerArgs { pub include_yul: bool, /// The output directory to write the output to or 'print' to print to the console - #[clap(long = "output", short = 'o', default_value = "")] + #[clap(long = "output", short = 'o', default_value = "output", hide_default_value = true)] pub output: String, } diff --git a/core/src/disassemble/mod.rs b/core/src/disassemble/mod.rs index 269988f1..b195b052 100644 --- a/core/src/disassemble/mod.rs +++ b/core/src/disassemble/mod.rs @@ -34,7 +34,7 @@ pub struct DisassemblerArgs { pub decimal_counter: bool, /// The output directory to write the output to or 'print' to print to the console - #[clap(long = "output", short = 'o', default_value = "output")] + #[clap(long = "output", short = 'o', default_value = "output", hide_default_value = true)] pub output: String, } diff --git a/core/src/dump/mod.rs b/core/src/dump/mod.rs index 7bd57aa9..8b136724 100644 --- a/core/src/dump/mod.rs +++ b/core/src/dump/mod.rs @@ -37,7 +37,7 @@ pub struct DumpArgs { pub verbose: clap_verbosity_flag::Verbosity, /// The output directory to write the output to or 'print' to print to the console - #[clap(long = "output", short, default_value = "", hide_default_value = true)] + #[clap(long = "output", short, default_value = "output", hide_default_value = true)] pub output: String, /// The RPC URL to use for fetching data. diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index fb31c249..cbe80bba 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -70,8 +70,8 @@ pub struct SnapshotArgs { #[clap(long)] pub no_tui: bool, - /// The output directory to write the output to - #[clap(long = "output", short = 'o', default_value = "")] + /// The output directory to write the output to, or 'print' to print to the console. + #[clap(long = "output", short = 'o', default_value = "output", hide_default_value = true)] pub output: String, } diff --git a/core/src/snapshot/util/csv.rs b/core/src/snapshot/util/csv.rs index c3b4e296..e5fa3b4e 100644 --- a/core/src/snapshot/util/csv.rs +++ b/core/src/snapshot/util/csv.rs @@ -2,18 +2,17 @@ use std::collections::HashMap; use heimdall_common::{ ether::signatures::{ResolvedError, ResolvedLog}, - utils::{io::file::write_lines_to_file, strings::encode_hex_reduced}, + utils::strings::encode_hex_reduced, }; use crate::snapshot::structures::snapshot::Snapshot; /// Write the snapshot data to a CSV file -pub fn generate_and_write_contract_csv( +pub fn generate_csv( snapshots: &Vec, resolved_errors: &HashMap, resolved_events: &HashMap, - output_path: &str, -) { +) -> Vec { let mut lines: Vec = Vec::new(); // add header @@ -132,5 +131,5 @@ pub fn generate_and_write_contract_csv( lines.push(line.join(",")); } - write_lines_to_file(output_path, lines); + lines } diff --git a/core/tests/test_cfg.rs b/core/tests/test_cfg.rs index 8c007929..115127e5 100644 --- a/core/tests/test_cfg.rs +++ b/core/tests/test_cfg.rs @@ -14,7 +14,6 @@ mod benchmark { rpc_url: String::from("https://eth.llamarpc.com"), default: true, color_edges: false, - format: String::from("png"), output: String::from(""), }; let _ = heimdall_core::cfg::cfg(args).await; @@ -32,7 +31,6 @@ mod benchmark { rpc_url: String::from("https://eth.llamarpc.com"), default: true, color_edges: false, - format: String::from("png"), output: String::from(""), }; let _ = heimdall_core::cfg::cfg(args).await; @@ -56,7 +54,6 @@ mod integration_tests { rpc_url: String::from("https://eth.llamarpc.com"), default: true, color_edges: false, - format: String::from("png"), output: String::from(""), }) .await @@ -80,7 +77,6 @@ mod integration_tests { rpc_url: String::from("https://eth.llamarpc.com"), default: true, color_edges: false, - format: String::from("png"), output: String::from(""), }) .await