chore: heimdall-common replacements (#241)

* refactor: replace duplicated code with get_bytecode_from_target * refactor: replace duplicated logger env set up with set_logger_env * refactor: replace duplicated target shortning with get_shortned_address * refactor: remove duplicated log env set up * style: code format * chore: replace logger setup with set_logger_env function for dump mod * chore: replace logger setup with set_logger_env function for inspect mod
Jon-Becker · Dec 20, 2023 · 769a33f · 769a33f
1 parent 20cbbe8
commit 769a33f
Show file tree

Hide file tree

Showing 7 changed files with 27 additions and 239 deletions.
diff --git a/core/src/cfg/mod.rs b/core/src/cfg/mod.rs
@@ -3,17 +3,16 @@ pub mod output;
 use derive_builder::Builder;
 use heimdall_common::{
     debug_max,
-    ether::{compiler::detect_compiler, rpc::get_code, selectors::find_function_selectors},
+    ether::{
+        bytecode::get_bytecode_from_target, compiler::detect_compiler,
+        selectors::find_function_selectors,
+    },
 };
 use indicatif::ProgressBar;
-use std::{fs, time::Duration};
+use std::time::Duration;
 
 use clap::{AppSettings, Parser};
-use heimdall_common::{
-    constants::{ADDRESS_REGEX, BYTECODE_REGEX},
-    ether::evm::core::vm::VM,
-    utils::io::logging::*,
-};
+use heimdall_common::{ether::evm::core::vm::VM, utils::io::logging::*};
 use petgraph::Graph;
 
 use crate::{
@@ -79,16 +78,7 @@ pub async fn cfg(args: CFGArgs) -> Result<Graph<String, String>, Box<dyn std::er
     use std::time::Instant;
     let now = Instant::now();
 
-    // set logger environment variable if not already set
-    if std::env::var("RUST_LOG").is_err() {
-        std::env::set_var(
-            "RUST_LOG",
-            match args.verbose.log_level() {
-                Some(level) => level.as_str(),
-                None => "SILENT",
-            },
-        );
-    }
+    set_logger_env(&args.verbose);
 
     let (logger, mut trace) = Logger::new(match args.verbose.log_level() {
         Some(level) => level.as_str(),
@@ -113,36 +103,7 @@ pub async fn cfg(args: CFGArgs) -> Result<Graph<String, String>, Box<dyn std::er
         "()".to_string(),
     );
 
-    // fetch bytecode
-    let contract_bytecode: String;
-    if ADDRESS_REGEX.is_match(&args.target).unwrap() {
-        // We are working with a contract address, so we need to fetch the bytecode from the RPC
-        // provider
-        contract_bytecode = get_code(&args.target, &args.rpc_url).await?;
-    } else if BYTECODE_REGEX.is_match(&args.target).unwrap() {
-        debug_max!("using provided bytecode for cfg generation");
-        contract_bytecode = args.target.replacen("0x", "", 1);
-    } else {
-        debug_max!("using provided file for cfg generation.");
-
-        // We are analyzing a file, so we need to read the bytecode from the file.
-        contract_bytecode = match fs::read_to_string(&args.target) {
-            Ok(contents) => {
-                let _contents = contents.replace('\n', "");
-                if BYTECODE_REGEX.is_match(&_contents).unwrap() && _contents.len() % 2 == 0 {
-                    _contents.replacen("0x", "", 1)
-                } else {
-                    logger
-                        .error(&format!("file '{}' doesn't contain valid bytecode.", &args.target));
-                    std::process::exit(1)
-                }
-            }
-            Err(_) => {
-                logger.error(&format!("failed to open file '{}' .", &args.target));
-                std::process::exit(1)
-            }
-        };
-    }
+    let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?;
 
     // disassemble the bytecode
     let disassembled_bytecode = disassemble(DisassemblerArgs {

diff --git a/core/src/decode/mod.rs b/core/src/decode/mod.rs
@@ -21,7 +21,10 @@ use heimdall_common::{
         signatures::{score_signature, ResolveSelector, ResolvedFunction},
     },
     utils::{
-        io::{logging::Logger, types::display},
+        io::{
+            logging::{set_logger_env, Logger},
+            types::display,
+        },
         strings::decode_hex,
     },
 };
@@ -94,16 +97,7 @@ impl DecodeArgsBuilder {
 /// calldata, without the ABI of the target contract.
 #[allow(deprecated)]
 pub async fn decode(args: DecodeArgs) -> Result<Vec<ResolvedFunction>, Error> {
-    // set logger environment variable if not already set
-    if std::env::var("RUST_LOG").is_err() {
-        std::env::set_var(
-            "RUST_LOG",
-            match args.verbose.log_level() {
-                Some(level) => level.as_str(),
-                None => "SILENT",
-            },
-        );
-    }
+    set_logger_env(&args.verbose);
 
     // get a new logger and trace
     let (logger, mut trace) = Logger::new(match args.verbose.log_level() {

diff --git a/core/src/decompile/mod.rs b/core/src/decompile/mod.rs
@@ -4,7 +4,9 @@ pub mod out;
 pub mod precompile;
 pub mod resolve;
 pub mod util;
-use heimdall_common::debug_max;
+use heimdall_common::{
+    debug_max, ether::bytecode::get_bytecode_from_target, utils::strings::get_shortned_target,
+};
 
 use crate::{
     decompile::{
@@ -20,17 +22,15 @@ use derive_builder::Builder;
 use heimdall_common::{
     ether::{
         compiler::detect_compiler,
-        rpc::get_code,
         selectors::{find_function_selectors, resolve_selectors},
     },
     utils::strings::encode_hex_reduced,
 };
 use indicatif::ProgressBar;
-use std::{collections::HashMap, fs, time::Duration};
+use std::{collections::HashMap, time::Duration};
 
 use clap::{AppSettings, Parser};
 use heimdall_common::{
-    constants::{ADDRESS_REGEX, BYTECODE_REGEX},
     ether::{evm::core::vm::VM, signatures::*},
     utils::io::logging::*,
 };
@@ -110,16 +110,7 @@ pub async fn decompile(
     use std::time::Instant;
     let now = Instant::now();
 
-    // set logger environment variable if not already set
-    if std::env::var("RUST_LOG").is_err() {
-        std::env::set_var(
-            "RUST_LOG",
-            match args.verbose.log_level() {
-                Some(level) => level.as_str(),
-                None => "SILENT",
-            },
-        );
-    }
+    set_logger_env(&args.verbose);
 
     // get a new logger
     let (logger, mut trace) = Logger::new(match args.verbose.log_level() {
@@ -136,13 +127,7 @@ pub async fn decompile(
         std::process::exit(1);
     }
 
-    // truncate target for prettier display
-    let mut shortened_target = args.target.clone();
-    if shortened_target.len() > 66 {
-        shortened_target = shortened_target.chars().take(66).collect::<String>() +
-            "..." +
-            &shortened_target.chars().skip(shortened_target.len() - 16).collect::<String>();
-    }
+    let shortened_target = get_shortned_target(&args.target);
     let decompile_call = trace.add_call(
         0,
         line!(),
@@ -152,37 +137,7 @@ pub async fn decompile(
         "()".to_string(),
     );
 
-    // parse the various formats that are accepted as targets
-    // i.e, file, bytecode, contract address
-    let contract_bytecode: String;
-    if ADDRESS_REGEX.is_match(&args.target)? {
-        // We are decompiling a contract address, so we need to fetch the bytecode from the RPC
-        // provider
-        contract_bytecode = get_code(&args.target, &args.rpc_url).await?;
-    } else if BYTECODE_REGEX.is_match(&args.target)? {
-        debug_max!("using provided bytecode for decompilation");
-        contract_bytecode = args.target.clone().replacen("0x", "", 1);
-    } else {
-        debug_max!("using provided file for decompilation.");
-
-        // We are decompiling a file, so we need to read the bytecode from the file.
-        contract_bytecode = match fs::read_to_string(&args.target) {
-            Ok(contents) => {
-                let _contents = contents.replace('\n', "");
-                if BYTECODE_REGEX.is_match(&_contents)? && _contents.len() % 2 == 0 {
-                    _contents.replacen("0x", "", 1)
-                } else {
-                    logger
-                        .error(&format!("file '{}' doesn't contain valid bytecode.", &args.target));
-                    std::process::exit(1)
-                }
-            }
-            Err(_) => {
-                logger.error(&format!("failed to open file '{}' .", &args.target));
-                std::process::exit(1)
-            }
-        };
-    }
+    let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?;
 
     // disassemble the bytecode
     let disassembled_bytecode = disassemble(DisassemblerArgs {

diff --git a/core/src/disassemble/mod.rs b/core/src/disassemble/mod.rs
@@ -1,12 +1,9 @@
-use std::fs;
-
 use clap::{AppSettings, Parser};
 use derive_builder::Builder;
 use heimdall_common::{
-    constants::{ADDRESS_REGEX, BYTECODE_REGEX},
-    ether::{evm::core::opcodes::Opcode, rpc::get_code},
+    ether::{bytecode::get_bytecode_from_target, evm::core::opcodes::Opcode},
     utils::{
-        io::logging::Logger,
+        io::logging::{set_logger_env, Logger},
         strings::{decode_hex, encode_hex},
     },
 };
@@ -60,56 +57,21 @@ pub async fn disassemble(args: DisassemblerArgs) -> Result<String, Box<dyn std::
     use std::time::Instant;
     let now = Instant::now();
 
-    // set logger environment variable if not already set
-    if std::env::var("RUST_LOG").is_err() {
-        std::env::set_var(
-            "RUST_LOG",
-            match args.verbose.log_level() {
-                Some(level) => level.as_str(),
-                None => "SILENT",
-            },
-        );
-    }
+    set_logger_env(&args.verbose);
 
     // get a new logger
     let (logger, _) = Logger::new(match args.verbose.log_level() {
         Some(level) => level.as_str(),
         None => "SILENT",
     });
 
-    let contract_bytecode: String;
-    if ADDRESS_REGEX.is_match(&args.target)? {
-        // We are disassembling a contract address, so we need to fetch the bytecode from the RPC
-        // provider.
-        contract_bytecode = get_code(&args.target, &args.rpc_url).await?;
-    } else if BYTECODE_REGEX.is_match(&args.target)? {
-        contract_bytecode = args.target;
-    } else {
-        // We are disassembling a file, so we need to read the bytecode from the file.
-        contract_bytecode = match fs::read_to_string(&args.target) {
-            Ok(contents) => {
-                let _contents = contents.replace('\n', "");
-                if BYTECODE_REGEX.is_match(&_contents)? && _contents.len() % 2 == 0 {
-                    _contents
-                } else {
-                    logger
-                        .error(&format!("file '{}' doesn't contain valid bytecode.", &args.target));
-                    std::process::exit(1)
-                }
-            }
-            Err(_) => {
-                logger.error(&format!("failed to open file '{}' .", &args.target));
-                std::process::exit(1)
-            }
-        };
-    }
+    let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?;
 
     let mut program_counter = 0;
     let mut output: String = String::new();
 
     // Iterate over the bytecode, disassembling each instruction.
     let byte_array = decode_hex(&contract_bytecode.replacen("0x", "", 1))?;
-
     while program_counter < byte_array.len() {
         let operation = Opcode::new(byte_array[program_counter]);
         let mut pushed_bytes: String = String::new();

diff --git a/core/src/dump/mod.rs b/core/src/dump/mod.rs
@@ -94,16 +94,7 @@ impl DumpArgsBuilder {
 /// entry point for the dump module. Will fetch all storage slots accessed by the target contract,
 /// and dump them to a CSV file or the TUI.
 pub async fn dump(args: DumpArgs) -> Result<Vec<DumpRow>, Box<dyn std::error::Error>> {
-    // set logger environment variable if not already set
-    if std::env::var("RUST_LOG").is_err() {
-        std::env::set_var(
-            "RUST_LOG",
-            match args.verbose.log_level() {
-                Some(level) => level.as_str(),
-                None => "SILENT",
-            },
-        );
-    }
+    set_logger_env(&args.verbose);
 
     let (logger, _) = Logger::new(match args.verbose.log_level() {
         Some(level) => level.as_str(),