Skip to content

Commit

Permalink
Refactor sierra-decompiler
Browse files Browse the repository at this point in the history
  • Loading branch information
@Rog3rSm1th
Rog3rSm1th committed Aug 20, 2024
1 parent 78b0534 commit 4fb9981
Show file tree
Hide file tree
Showing 2 changed files with 156 additions and 122 deletions.
2 changes: 1 addition & 1 deletion lib/src/lib.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
mod config;
mod decompiler;
pub mod decompiler;
pub mod detectors;
pub mod graph;
pub mod provider;
Expand Down
276 changes: 155 additions & 121 deletions sierra-decompiler/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@ use serde_json;
use tokio;

use cairo_lang_starknet_classes::contract_class::ContractClass;
use sierra_analyzer_lib::decompiler::decompiler::Decompiler;
use sierra_analyzer_lib::detectors::get_detectors;
use sierra_analyzer_lib::graph::graph::save_svg_graph_to_file;
use sierra_analyzer_lib::provider::NetworkConfig;
use sierra_analyzer_lib::provider::RpcClient;
use sierra_analyzer_lib::sierra_program;
use sierra_analyzer_lib::sierra_program::SierraProgram;

/// Decompile a Sierra program
Expand Down Expand Up @@ -68,67 +68,23 @@ struct Args {
async fn main() {
let args = Args::parse();

// Ensure either remote or Sierra file is provided
if args.remote.is_empty() && args.sierra_file.is_none() {
eprintln!("Error: Either remote or Sierra file must be provided");
return;
}

// Define program and sierra_file before the if statement
let program: SierraProgram;
let mut sierra_file: Option<PathBuf> = None;

// Analyze a contract deployed on Starknet
if !args.remote.is_empty() {
// Define the client based on the network parameter
let client = match args.network.as_str() {
"mainnet" => RpcClient::new(NetworkConfig::MAINNET_API_URL),
"sepolia" => RpcClient::new(NetworkConfig::SEPOLIA_API_URL),
_ => {
eprintln!("Error: Unsupported network type '{}'", args.network);
return;
}
};

// Fetch contract class from the RPC Node
match client.get_class(&args.remote).await {
Ok(response) => {
// Convert RpcClient response to JSON content
let content = response.to_json();

// Deserialize JSON into a ContractClass
let program_string = serde_json::from_str::<ContractClass>(&content)
.ok()
.and_then(|prog| prog.extract_sierra_program().ok())
.map_or_else(|| content.clone(), |prog_sierra| prog_sierra.to_string());
program = SierraProgram::new(program_string);
}
Err(e) => {
eprintln!("Error calling RPC: {}", e);
// Stop the program if there is an error in the RPC response
return;
}
// Load the Sierra program
let program = match load_program(&args).await {
Ok(program) => program,
Err(e) => {
eprintln!("Error loading program: {}", e);
return;
}
}
// Analyze a local file
else {
sierra_file = args.sierra_file;
let mut file = File::open(sierra_file.as_ref().unwrap()).expect("Failed to open file");
let mut content = String::new();
file.read_to_string(&mut content)
.expect("Failed to read file");

// Deserialize JSON into a ContractClass, or use the content directly if that fails
let program_string = serde_json::from_str::<ContractClass>(&content)
.ok()
.and_then(|prog| prog.extract_sierra_program().ok())
.map_or_else(|| content.clone(), |prog_sierra| prog_sierra.to_string());
program = sierra_program::SierraProgram::new(program_string);
}
};

// Color output by default and if CFG or Callgraph is not enabled to avoid bugs in the SVG output
// Determine if colored output is needed
let colored_output = !args.no_color ^ (args.cfg | args.callgraph);

// Now you can use program and sierra_file outside the if and else blocks
let mut decompiler = program.decompiler(args.verbose);
let decompiled_code = decompiler.decompile(colored_output);

Expand All @@ -137,84 +93,162 @@ async fn main() {
decompiler.filter_functions(function_name);
}

// Determine the file stem based on the remote address or the sierra_file
let file_stem = if !args.remote.is_empty() {
// Determine the file stem based on the remote address or the Sierra file
let file_stem = get_file_stem(&args);

// Handle different output options
// CFG
if args.cfg {
handle_cfg(&args, &mut decompiler, &file_stem);
}
// Callgraph
else if args.callgraph {
handle_callgraph(&args, &mut decompiler, &file_stem);
}
// Detectors
else if args.detectors {
handle_detectors(&mut decompiler);
}
// Decompiler (default)
else {
println!("{}", decompiled_code);
}
}

/// Load the Sierra program from either a remote source or a local file
async fn load_program(args: &Args) -> Result<SierraProgram, String> {
if !args.remote.is_empty() {
load_remote_program(args).await
} else {
load_local_program(args)
}
}

/// Load the Sierra program from a remote source
async fn load_remote_program(args: &Args) -> Result<SierraProgram, String> {
let client = match args.network.as_str() {
"mainnet" => RpcClient::new(NetworkConfig::MAINNET_API_URL),
"sepolia" => RpcClient::new(NetworkConfig::SEPOLIA_API_URL),
_ => {
return Err(format!(
"Error: Unsupported network type '{}'",
args.network
))
}
};

match client.get_class(&args.remote).await {
Ok(response) => {
let content = response.to_json();
let program_string = serde_json::from_str::<ContractClass>(&content)
.ok()
.and_then(|prog| prog.extract_sierra_program().ok())
.map_or_else(|| content.clone(), |prog_sierra| prog_sierra.to_string());
Ok(SierraProgram::new(program_string))
}
Err(e) => Err(format!("Error calling RPC: {}", e)),
}
}

/// Load the Sierra program from a local file
fn load_local_program(args: &Args) -> Result<SierraProgram, String> {
let sierra_file = args.sierra_file.as_ref().unwrap();
let mut file = File::open(sierra_file).map_err(|e| format!("Failed to open file: {}", e))?;
let mut content = String::new();
file.read_to_string(&mut content)
.map_err(|e| format!("Failed to read file: {}", e))?;

let program_string = serde_json::from_str::<ContractClass>(&content)
.ok()
.and_then(|prog| prog.extract_sierra_program().ok())
.map_or_else(|| content.clone(), |prog_sierra| prog_sierra.to_string());
Ok(SierraProgram::new(program_string))
}

/// Get the file stem based on the remote address or the Sierra file
fn get_file_stem(args: &Args) -> String {
if !args.remote.is_empty() {
args.remote.clone()
} else {
sierra_file
args.sierra_file
.as_ref()
.unwrap()
.file_stem()
.unwrap_or_default()
.to_string_lossy()
.to_string()
};
}
}

if args.cfg {
let svg_filename = format!("{}_cfg.svg", file_stem);
let full_path = args.cfg_output.join(svg_filename);

// Create the output directory if it doesn't exist
if let Err(e) = fs::create_dir_all(&args.cfg_output) {
eprintln!(
"Failed to create directory '{}': {}",
args.cfg_output.display(),
e
);
return;
}
/// Handle the generation and saving of the CFG (Control Flow Graph)
fn handle_cfg(args: &Args, decompiler: &mut Decompiler, file_stem: &str) {
let svg_filename = format!("{}_cfg.svg", file_stem);
let full_path = args.cfg_output.join(svg_filename);

// Create the output directory if it doesn't exist
if let Err(e) = fs::create_dir_all(&args.cfg_output) {
eprintln!(
"Failed to create directory '{}': {}",
args.cfg_output.display(),
e
);
return;
}

// Generate CFG and save to SVG
let cfg_graph = decompiler.generate_cfg();
save_svg_graph_to_file(full_path.to_str().unwrap(), cfg_graph)
.expect("Failed to save CFG to SVG");
} else if args.callgraph {
let svg_filename = format!("{}_callgraph.svg", file_stem);
let full_path = args.callgraph_output.join(svg_filename);

// Create the output directory if it doesn't exist
if let Err(e) = fs::create_dir_all(&args.callgraph_output) {
eprintln!(
"Failed to create directory '{}': {}",
args.callgraph_output.display(),
e
);
return;
}
// Generate CFG and save to SVG
let cfg_graph = decompiler.generate_cfg();
save_svg_graph_to_file(full_path.to_str().unwrap(), cfg_graph)
.expect("Failed to save CFG to SVG");
}

// Generate Callgraph and save to SVG
let callgraph_graph = decompiler.generate_callgraph();
save_svg_graph_to_file(full_path.to_str().unwrap(), callgraph_graph)
.expect("Failed to save Callgraph to SVG");
} else if args.detectors {
let mut detectors = get_detectors();
let mut output = String::new();

// Run all the detectors
for detector in detectors.iter_mut() {
let result = detector.detect(&mut decompiler);
if !result.trim().is_empty() {
// Each detector output is formatted like
//
// [Detector category] Detector name
// - detector content
// - ...
output.push_str(&format!(
"[{}] {}\n{}\n\n",
detector.detector_type().as_str(),
detector.name(),
result
.lines()
.map(|line| format!("\t- {}", line))
.collect::<Vec<String>>()
.join("\n")
));
}
}
/// Handle the generation and saving of the Call Graph
fn handle_callgraph(args: &Args, decompiler: &mut Decompiler, file_stem: &str) {
let svg_filename = format!("{}_callgraph.svg", file_stem);
let full_path = args.callgraph_output.join(svg_filename);

// Create the output directory if it doesn't exist
if let Err(e) = fs::create_dir_all(&args.callgraph_output) {
eprintln!(
"Failed to create directory '{}': {}",
args.callgraph_output.display(),
e
);
return;
}

// Print the detectors result
println!("{}", output.trim());
} else {
println!("{}", decompiled_code);
// Generate Callgraph and save to SVG
let callgraph_graph = decompiler.generate_callgraph();
save_svg_graph_to_file(full_path.to_str().unwrap(), callgraph_graph)
.expect("Failed to save Callgraph to SVG");
}

/// Handle the running of detectors and printing their results
fn handle_detectors(decompiler: &mut Decompiler) {
let mut detectors = get_detectors();
let mut output = String::new();

// Run all the detectors
for detector in detectors.iter_mut() {
let result = detector.detect(decompiler);
if !result.trim().is_empty() {
// Each detector output is formatted like
//
// [Detector category] Detector name
// - detector content
// - ...
output.push_str(&format!(
"[{}] {}\n{}\n\n",
detector.detector_type().as_str(),
detector.name(),
result
.lines()
.map(|line| format!("\t- {}", line))
.collect::<Vec<String>>()
.join("\n")
));
}
}

// Print the detectors result
println!("{}", output.trim());
}

0 comments on commit 4fb9981

Please sign in to comment.