Skip to content

Commit

Permalink
fix(cli): standardize --output flag paths, opts
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jon-Becker
Jon-Becker committed Nov 23, 2023
1 parent 0eafb0e commit 8950f3f
Show file tree
Hide file tree
Showing 10 changed files with 147 additions and 196 deletions.
163 changes: 59 additions & 104 deletions cli/src/main.rs
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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};

Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -99,15 +98,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}));

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;
Expand All @@ -116,17 +108,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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);
}
}

Expand All @@ -138,43 +127,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {

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()
Expand All @@ -196,13 +164,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
),
);
}

// 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);
}
}
}
Expand Down Expand Up @@ -231,20 +203,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}

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) => {
Expand All @@ -261,9 +228,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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"));

Expand All @@ -280,16 +244,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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);
}
}
Expand All @@ -301,27 +258,25 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}

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);
}
Expand Down
76 changes: 76 additions & 0 deletions cli/src/output.rs
View file
Original file line number Diff line number Diff line change
@@ -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<String, Box<dyn std::error::Error>> {
// 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());
}
}
8 changes: 1 addition & 7 deletions core/src/cfg/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand All @@ -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()),
}
Expand Down
Loading

0 comments on commit 8950f3f

Please sign in to comment.