Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(cli): add --output flag #189

Merged
merged 12 commits into from
Nov 23, 2023
Merged

feat(cli): add --output flag #189

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
9b32275
decompile output option init
solidoracle Nov 20, 2023
1c0664a
fix path creation
solidoracle Nov 21, 2023
40aefd5
Merge remote-tracking branch 'upstream/main' into output-option
solidoracle Nov 21, 2023
e61f684
output default_value
solidoracle Nov 21, 2023
665c551
Decompile output option
solidoracle Nov 21, 2023
c7d0ec1
CFG output option
solidoracle Nov 21, 2023
07baf18
Dump output option implementation
solidoracle Nov 21, 2023
3e634c3
fix output comment
solidoracle Nov 21, 2023
e8d9f78
Snapshot output option
solidoracle Nov 22, 2023
7df4bc4
fix indenting
solidoracle Nov 22, 2023
0eafb0e
chore(cli): fmt & clippy
Jon-Becker Nov 22, 2023
8950f3f
fix(cli): standardize `--output` flag paths, opts
Jon-Becker Nov 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 99 additions & 114 deletions cli/src/main.rs
View file
Open in desktop
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 All @@ -77,7 +76,6 @@ pub enum Subcommands {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = Arguments::parse();

// handle catching panics with
panic::set_hook(Box::new(|panic_info| {
// cleanup the terminal (break out of alternate screen, disable mouse capture, and show the
Expand All @@ -100,11 +98,6 @@ 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();
output_path.push_str("/output");

match args.sub {
Subcommands::Disassemble(mut cmd) => {
// if the user has not specified a rpc url, use the default
Expand All @@ -114,17 +107,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {

let assembly = disassemble(cmd.clone()).await?;

// write to file
if ADDRESS_REGEX.is_match(&cmd.target).unwrap() {
output_path.push_str(&format!(
"/{}/{}/disassembled.asm",
rpc::chain_id(&cmd.rpc_url).await.unwrap(),
&cmd.target
));
if cmd.output == "print" {
// TODO: use `less`
println!("{}", assembly);
} else {
output_path.push_str("/local/disassembled.asm");
let output_path =
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "disassembled.asm")
.await?;

write_file(&output_path, &assembly);
}
write_file(&output_path, &assembly);
}

Subcommands::Decompile(mut cmd) => {
Expand All @@ -135,55 +127,54 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {

let result = decompile(cmd.clone()).await?;

// write to file
let abi_output_path;
let solidity_output_path;
let yul_output_path;
if ADDRESS_REGEX.is_match(&cmd.target).unwrap() {
let chain_id = rpc::chain_id(&cmd.rpc_url).await.unwrap();

abi_output_path =
format!("{}/{}/{}/abi.json", &output_path, &chain_id, &cmd.target);
solidity_output_path =
format!("{}/{}/{}/decompiled.sol", &output_path, &chain_id, &cmd.target);
yul_output_path =
format!("{}/{}/{}/decompiled.yul", &output_path, &chain_id, &cmd.target);
if cmd.output == "print" {
if let Some(abi) = &result.abi {
println!("ABI:\n\n{}\n", serde_json::to_string_pretty(abi).unwrap());
}
if let Some(source) = &result.source {
println!("Source:\n\n{}\n", source);
}
} else {
abi_output_path = format!("{}/local/abi.json", &output_path);
solidity_output_path = format!("{}/local/decompiled.sol", &output_path);
yul_output_path = format!("{}/local/decompiled.yul", &output_path);
}

if let Some(abi) = result.abi {
// write the ABI to a file
write_file(
&abi_output_path,
&format!(
"[{}]",
abi.iter()
.map(|x| {
match x {
ABIStructure::Function(x) => {
serde_json::to_string_pretty(x).unwrap()
}
ABIStructure::Error(x) => {
serde_json::to_string_pretty(x).unwrap()
}
ABIStructure::Event(x) => {
serde_json::to_string_pretty(x).unwrap()
// write the contract ABI
if let Some(abi) = result.abi {
let output_path =
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "abi.json")
.await?;

write_file(
&output_path,
&format!(
"[{}]",
abi.iter()
.map(|x| {
match x {
ABIStructure::Function(x) => {
serde_json::to_string_pretty(x).unwrap()
}
ABIStructure::Error(x) => {
serde_json::to_string_pretty(x).unwrap()
}
ABIStructure::Event(x) => {
serde_json::to_string_pretty(x).unwrap()
}
}
}
})
.collect::<Vec<String>>()
.join(",\n")
),
);
}
if let Some(source) = result.source {
if cmd.include_solidity {
write_file(&solidity_output_path, &source);
} else {
write_file(&yul_output_path, &source);
})
.collect::<Vec<String>>()
.join(",\n")
),
);
}

// write the contract source
if let Some(source) = &result.source {
let output_path = if cmd.include_solidity {
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "decompiled.sol")
.await?
} else {
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "decompiled.yul")
.await?
};
write_file(&output_path, source);
}
}
}
Expand Down Expand Up @@ -212,19 +203,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}

let cfg = cfg(cmd.clone()).await?;
let stringified_dot = build_cfg(&cfg, &cmd);

// write to file
if ADDRESS_REGEX.is_match(&cmd.target).unwrap() {
output_path.push_str(&format!(
"/{}/{}",
rpc::chain_id(&cmd.rpc_url).await.unwrap(),
&cmd.target
));
if cmd.output == "print" {
println!("{}", stringified_dot);
} else {
output_path.push_str("/local");
let output_path =
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "cfg.dot").await?;
write_file(&output_path, &stringified_dot);
}

write_cfg_to_file(&cfg, &cmd, output_path)
}

Subcommands::Dump(mut cmd) => {
Expand All @@ -241,17 +228,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let result = dump(cmd.clone()).await?;
let mut lines = Vec::new();

// 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");
}

// add header
lines.push(String::from("last_modified,alias,slot,decoded_type,value"));

Expand All @@ -263,8 +239,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
));
}

// write to file
write_lines_to_file(&output_path, lines);
if cmd.output == "print" {
for line in &lines {
println!("{}", line);
}
} else {
let output_path =
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "dump.csv").await?;

write_lines_to_file(&output_path, lines);
}
}

Subcommands::Snapshot(mut cmd) => {
Expand All @@ -273,25 +257,26 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
cmd.rpc_url = configuration.rpc_url;
}

// write to file
if ADDRESS_REGEX.is_match(&cmd.target).unwrap() {
output_path.push_str(&format!(
"/{}/{}/snapshot.csv",
rpc::chain_id(&cmd.rpc_url).await.unwrap(),
&cmd.target,
));
let snapshot_result = snapshot(cmd.clone()).await?;
let csv_lines = generate_csv(
&snapshot_result.snapshots,
&snapshot_result.resolved_errors,
&snapshot_result.resolved_events,
);

if cmd.output == "print" {
for line in &csv_lines {
println!("{}", line);
}
} else {
output_path.push_str("/local/snapshot.csv");
}
let output_path =
build_output_path(&cmd.output, &cmd.target, &cmd.rpc_url, "snapshot.csv")
.await?;

let snapshot = snapshot(cmd).await?;
generate_and_write_contract_csv(
&snapshot.snapshots,
&snapshot.resolved_errors,
&snapshot.resolved_events,
&output_path,
)
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
Open in desktop
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());
}
}
Loading
Loading