diff --git a/Cargo.lock b/Cargo.lock index 56e28286..f3c37df7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1934,6 +1934,7 @@ version = "0.8.3" dependencies = [ "alloy", "alloy-json-abi", + "async-trait", "clap", "colored", "eyre", diff --git a/crates/cfg/src/core/mod.rs b/crates/cfg/src/core/mod.rs index 5dec69d9..9ac86621 100644 --- a/crates/cfg/src/core/mod.rs +++ b/crates/cfg/src/core/mod.rs @@ -2,10 +2,7 @@ pub(crate) mod graph; use alloy::primitives::Address; use eyre::eyre; -use heimdall_common::{ - ether::{bytecode::get_bytecode_from_target, compiler::detect_compiler}, - utils::strings::StringExt, -}; +use heimdall_common::{ether::compiler::detect_compiler, utils::strings::StringExt}; use heimdall_vm::core::vm::VM; use petgraph::{dot::Dot, Graph}; @@ -53,7 +50,8 @@ pub async fn cfg(args: CFGArgs) -> Result { // get the bytecode from the target let start_fetch_time = Instant::now(); - let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url) + let contract_bytecode = args + .get_bytecode() .await .map_err(|e| Error::FetchError(format!("fetching target bytecode failed: {}", e)))?; debug!("fetching target bytecode took {:?}", start_fetch_time.elapsed()); diff --git a/crates/cfg/src/interfaces/args.rs b/crates/cfg/src/interfaces/args.rs index 3c3a8a23..3338ad5a 100644 --- a/crates/cfg/src/interfaces/args.rs +++ b/crates/cfg/src/interfaces/args.rs @@ -1,5 +1,7 @@ use clap::Parser; use derive_builder::Builder; +use eyre::Result; +use heimdall_common::ether::bytecode::get_bytecode_from_target; use heimdall_config::parse_url_arg; #[derive(Debug, Clone, Parser, Builder)] @@ -40,6 +42,12 @@ pub struct CFGArgs { pub timeout: u64, } +impl CFGArgs { + pub async fn get_bytecode(&self) -> Result> { + get_bytecode_from_target(&self.target, &self.rpc_url).await + } +} + impl CFGArgsBuilder { pub fn new() -> Self { Self { diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 9c36bb93..0be74f55 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -27,6 +27,7 @@ tracing-subscriber = "0.3.18" eyre = "0.6.12" alloy-json-abi = "0.7.6" alloy = { version = "0.1.3", features = ["full", "rpc-types-debug", "rpc-types-trace"] } +async-trait = "0.1.51" [[bin]] name = "heimdall" diff --git a/crates/cli/src/log_args.rs b/crates/cli/src/args.rs similarity index 70% rename from crates/cli/src/log_args.rs rename to crates/cli/src/args.rs index b0c6a34d..248283b6 100644 --- a/crates/cli/src/log_args.rs +++ b/crates/cli/src/args.rs @@ -1,7 +1,13 @@ -//! clap [Args](clap::Args) for logging configuration. -// Mostly taken from [reth](https://github.com/paradigmxyz/reth) +use clap::{Parser, Subcommand}; use clap::{ArgAction, Args, ValueEnum}; +use heimdall_cache::CacheArgs; +use heimdall_config::ConfigArgs; +use heimdall_core::{ + heimdall_cfg::CFGArgs, heimdall_decoder::DecodeArgs, heimdall_decompiler::DecompilerArgs, + heimdall_disassembler::DisassemblerArgs, heimdall_dump::DumpArgs, + heimdall_inspect::InspectArgs, +}; use heimdall_tracing::{ tracing_subscriber::filter::Directive, FileWorkerGuard, HeimdallTracer, LayerInfo, LogFormat, Tracer, @@ -12,6 +18,51 @@ use std::{ }; use tracing::{level_filters::LevelFilter, Level}; +#[derive(Debug, Parser)] +#[clap(name = "heimdall", author = "Jonathan Becker ", version)] +pub struct Arguments { + #[clap(subcommand)] + pub sub: Subcommands, + + #[clap(flatten)] + pub logs: LogArgs, +} + +#[derive(Debug, Subcommand)] +#[clap( + about = "Heimdall is an advanced Ethereum smart contract toolkit for forensic and heuristic analysis.", + after_help = "For more information, read the wiki: https://jbecker.dev/r/heimdall-rs/wiki" +)] +#[allow(clippy::large_enum_variant)] +pub enum Subcommands { + #[clap(name = "disassemble", about = "Disassemble EVM bytecode to assembly")] + Disassemble(DisassemblerArgs), + + #[clap(name = "decompile", about = "Decompile EVM bytecode to Solidity")] + Decompile(DecompilerArgs), + + #[clap(name = "cfg", about = "Generate a visual control flow graph for EVM bytecode")] + CFG(CFGArgs), + + #[clap(name = "decode", about = "Decode calldata into readable types")] + Decode(DecodeArgs), + + #[clap(name = "config", about = "Display and edit the current configuration")] + Config(ConfigArgs), + + #[clap(name = "cache", about = "Manage heimdall-rs' cached files")] + Cache(CacheArgs), + + #[clap(name = "dump", about = "Dump the value of all storage slots accessed by a contract")] + Dump(DumpArgs), + + #[clap( + name = "inspect", + about = "Detailed inspection of Ethereum transactions, including calldata & trace decoding, log visualization, and more" + )] + Inspect(InspectArgs), +} + /// The log configuration. #[derive(Debug, Args)] #[clap(next_help_heading = "LOGGING")] diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 1fd54f71..942cafcf 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -1,74 +1,24 @@ -pub(crate) mod log_args; +pub(crate) mod args; pub(crate) mod output; +use args::{Arguments, Subcommands}; +use clap::Parser; use eyre::{eyre, Result}; -use log_args::LogArgs; +use heimdall_cache::cache; use output::{build_output_path, print_with_less}; use tracing::info; -use clap::{Parser, Subcommand}; - -use heimdall_cache::{cache, CacheArgs}; use heimdall_common::utils::{ hex::ToLowerHex, io::file::write_file, version::{current_version, remote_nightly_version, remote_version}, }; -use heimdall_config::{config, ConfigArgs, Configuration}; +use heimdall_config::{config, Configuration}; use heimdall_core::{ - heimdall_cfg::{cfg, CFGArgs}, - heimdall_decoder::{decode, DecodeArgs}, - heimdall_decompiler::{decompile, DecompilerArgs}, - heimdall_disassembler::{disassemble, DisassemblerArgs}, - heimdall_dump::{dump, DumpArgs}, - heimdall_inspect::{inspect, InspectArgs}, + heimdall_cfg::cfg, heimdall_decoder::decode, heimdall_decompiler::decompile, + heimdall_disassembler::disassemble, heimdall_dump::dump, heimdall_inspect::inspect, }; -#[derive(Debug, Parser)] -#[clap(name = "heimdall", author = "Jonathan Becker ", version)] -pub struct Arguments { - #[clap(subcommand)] - pub sub: Subcommands, - - #[clap(flatten)] - logs: LogArgs, -} - -#[derive(Debug, Subcommand)] -#[clap( - about = "Heimdall is an advanced Ethereum smart contract toolkit for forensic and heuristic analysis.", - after_help = "For more information, read the wiki: https://jbecker.dev/r/heimdall-rs/wiki" -)] -#[allow(clippy::large_enum_variant)] -pub enum Subcommands { - #[clap(name = "disassemble", about = "Disassemble EVM bytecode to assembly")] - Disassemble(DisassemblerArgs), - - #[clap(name = "decompile", about = "Decompile EVM bytecode to Solidity")] - Decompile(DecompilerArgs), - - #[clap(name = "cfg", about = "Generate a visual control flow graph for EVM bytecode")] - CFG(CFGArgs), - - #[clap(name = "decode", about = "Decode calldata into readable types")] - Decode(DecodeArgs), - - #[clap(name = "config", about = "Display and edit the current configuration")] - Config(ConfigArgs), - - #[clap(name = "cache", about = "Manage heimdall-rs' cached files")] - Cache(CacheArgs), - - #[clap(name = "dump", about = "Dump the value of all storage slots accessed by a contract")] - Dump(DumpArgs), - - #[clap( - name = "inspect", - about = "Detailed inspection of Ethereum transactions, including calldata & trace decoding, log visualization, and more" - )] - Inspect(InspectArgs), -} - #[tokio::main] async fn main() -> Result<()> { let args = Arguments::parse(); diff --git a/crates/common/src/ether/bytecode.rs b/crates/common/src/ether/bytecode.rs index 0d5fa7a8..a95d158d 100644 --- a/crates/common/src/ether/bytecode.rs +++ b/crates/common/src/ether/bytecode.rs @@ -5,9 +5,7 @@ use alloy::primitives::{bytes::Bytes, Address}; use eyre::{eyre, Result}; use std::fs; -/// Given a target from the CLI, return bytecode of the target. -/// TODO: this can probably be a trait method so we can do something like target.try_get_bytecode() -/// TODO: move to CLI, since its only used in CLI +/// Given a target, return bytecode of the target. pub async fn get_bytecode_from_target(target: &str, rpc_url: &str) -> Result> { // If the target is an address, fetch the bytecode from the RPC provider. if let Ok(address) = target.parse::
() { diff --git a/crates/common/src/ether/calldata.rs b/crates/common/src/ether/calldata.rs index 07e1c6f7..80cc2cf3 100644 --- a/crates/common/src/ether/calldata.rs +++ b/crates/common/src/ether/calldata.rs @@ -3,7 +3,7 @@ use crate::utils::strings::decode_hex; use alloy::primitives::TxHash; use eyre::{bail, eyre, Result}; -/// Given a target from the CLI, return calldata of the target. +/// Given a target, return calldata of the target. pub async fn get_calldata_from_target(target: &str, rpc_url: &str) -> Result> { // If the target is a transaction hash, fetch the calldata from the RPC provider. if let Ok(address) = target.parse::() { diff --git a/crates/decode/src/core/mod.rs b/crates/decode/src/core/mod.rs index 3548487e..ac87a043 100644 --- a/crates/decode/src/core/mod.rs +++ b/crates/decode/src/core/mod.rs @@ -5,7 +5,6 @@ use alloy_dyn_abi::{DynSolCall, DynSolReturns, DynSolType}; use eyre::eyre; use heimdall_common::{ ether::{ - calldata::get_calldata_from_target, signatures::{score_signature, ResolveSelector, ResolvedFunction}, types::parse_function_parameters, }, @@ -44,7 +43,8 @@ pub async fn decode(mut args: DecodeArgs) -> Result { // get the bytecode from the target let start_fetch_time = Instant::now(); - let mut calldata = get_calldata_from_target(&args.target, &args.rpc_url) + let mut calldata = args + .get_calldata() .await .map_err(|e| Error::FetchError(format!("fetching target calldata failed: {}", e)))?; debug!("fetching target calldata took {:?}", start_fetch_time.elapsed()); diff --git a/crates/decode/src/interfaces/args.rs b/crates/decode/src/interfaces/args.rs index 2e04752e..e13c05bd 100644 --- a/crates/decode/src/interfaces/args.rs +++ b/crates/decode/src/interfaces/args.rs @@ -1,5 +1,7 @@ use clap::Parser; use derive_builder::Builder; +use eyre::Result; +use heimdall_common::ether::calldata::get_calldata_from_target; use heimdall_config::parse_url_arg; #[derive(Debug, Clone, Parser, Builder)] @@ -43,6 +45,12 @@ pub struct DecodeArgs { pub skip_resolving: bool, } +impl DecodeArgs { + pub async fn get_calldata(&self) -> Result> { + get_calldata_from_target(&self.target, &self.rpc_url).await + } +} + impl DecodeArgsBuilder { pub fn new() -> Self { Self { diff --git a/crates/decompile/src/core/mod.rs b/crates/decompile/src/core/mod.rs index 077fe0b6..9bd98ad5 100644 --- a/crates/decompile/src/core/mod.rs +++ b/crates/decompile/src/core/mod.rs @@ -9,7 +9,6 @@ use alloy_json_abi::JsonAbi; use eyre::eyre; use heimdall_common::{ ether::{ - bytecode::get_bytecode_from_target, compiler::detect_compiler, signatures::{score_signature, ResolvedError, ResolvedFunction, ResolvedLog}, types::to_type, @@ -60,7 +59,8 @@ pub async fn decompile(args: DecompilerArgs) -> Result { // get the bytecode from the target let start_fetch_time = Instant::now(); - let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url) + let contract_bytecode = args + .get_bytecode() .await .map_err(|e| Error::FetchError(format!("fetching target bytecode failed: {}", e)))?; debug!("fetching target bytecode took {:?}", start_fetch_time.elapsed()); diff --git a/crates/decompile/src/interfaces/args.rs b/crates/decompile/src/interfaces/args.rs index 2f879537..47f1ea43 100644 --- a/crates/decompile/src/interfaces/args.rs +++ b/crates/decompile/src/interfaces/args.rs @@ -1,5 +1,7 @@ use clap::Parser; use derive_builder::Builder; +use eyre::Result; +use heimdall_common::ether::bytecode::get_bytecode_from_target; use heimdall_config::parse_url_arg; #[derive(Debug, Clone, Parser, Builder)] @@ -47,6 +49,12 @@ pub struct DecompilerArgs { pub timeout: u64, } +impl DecompilerArgs { + pub async fn get_bytecode(&self) -> Result> { + get_bytecode_from_target(&self.target, &self.rpc_url).await + } +} + impl DecompilerArgsBuilder { pub fn new() -> Self { Self { diff --git a/crates/disassemble/src/core/mod.rs b/crates/disassemble/src/core/mod.rs index 107e82e0..eeb125e5 100644 --- a/crates/disassemble/src/core/mod.rs +++ b/crates/disassemble/src/core/mod.rs @@ -2,7 +2,7 @@ use std::time::Instant; use crate::{error::Error, interfaces::DisassemblerArgs}; use eyre::eyre; -use heimdall_common::{ether::bytecode::get_bytecode_from_target, utils::strings::encode_hex}; +use heimdall_common::utils::strings::encode_hex; use heimdall_vm::core::opcodes::Opcode; use tracing::{debug, info}; @@ -14,9 +14,8 @@ pub async fn disassemble(args: DisassemblerArgs) -> Result { // get the bytecode from the target let start_fetch_time = Instant::now(); - let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url) - .await - .map_err(|e| eyre!("fetching target bytecode failed: {}", e))?; + let contract_bytecode = + args.get_bytecode().await.map_err(|e| eyre!("fetching target bytecode failed: {}", e))?; debug!("fetching target bytecode took {:?}", start_fetch_time.elapsed()); // iterate over the bytecode, disassembling each instruction diff --git a/crates/disassemble/src/interfaces/args.rs b/crates/disassemble/src/interfaces/args.rs index 76c4ff1a..713b3446 100644 --- a/crates/disassemble/src/interfaces/args.rs +++ b/crates/disassemble/src/interfaces/args.rs @@ -1,4 +1,6 @@ use clap::Parser; +use eyre::Result; +use heimdall_common::ether::bytecode::get_bytecode_from_target; use heimdall_config::parse_url_arg; #[derive(Debug, Clone, Parser)] @@ -48,6 +50,12 @@ pub struct DisassemblerArgsBuilder { output: Option, } +impl DisassemblerArgs { + pub async fn get_bytecode(&self) -> Result> { + get_bytecode_from_target(&self.target, &self.rpc_url).await + } +} + impl Default for DisassemblerArgsBuilder { fn default() -> Self { Self::new()