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

chore: heimdall-common replacements #241

Merged
merged 7 commits into from
Dec 20, 2023
Merged

chore: heimdall-common replacements #241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ae15076
refactor: replace duplicated code with get_bytecode_from_target
iankressin Dec 14, 2023
e0203bb
refactor: replace duplicated logger env set up with set_logger_env
iankressin Dec 14, 2023
c9ca161
refactor: replace duplicated target shortning with get_shortned_address
iankressin Dec 14, 2023
24ce191
refactor: remove duplicated log env set up
iankressin Dec 15, 2023
53e7224
style: code format
iankressin Dec 15, 2023
84826dc
chore: replace logger setup with set_logger_env function for dump mod
iankressin Dec 18, 2023
e9bdcd5
chore: replace logger setup with set_logger_env function for inspect mod
iankressin Dec 18, 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
55 changes: 8 additions & 47 deletions core/src/cfg/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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(),
Expand All @@ -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 {
Expand Down
16 changes: 5 additions & 11 deletions core/src/decode/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
};
Expand Down Expand Up @@ -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() {
Expand Down
59 changes: 7 additions & 52 deletions core/src/decompile/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand All @@ -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::*,
};
Expand Down Expand Up @@ -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() {
Expand All @@ -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!(),
Expand All @@ -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 {
Expand Down
46 changes: 4 additions & 42 deletions core/src/disassemble/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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},
},
};
Expand Down Expand Up @@ -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();
Expand Down
11 changes: 1 addition & 10 deletions core/src/dump/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
Loading
Loading