use crate::processor::{ new_processing_state_background_thread, CompressionTypeField, FileDataKindField, IncludedFunctionsField, SaveReportToDiskField, WarpFileProcessor, }; use crate::report::{ReportGenerator, ReportKindField}; use crate::{user_signature_dir, INCLUDE_TAG_NAME}; use binaryninja::background_task::BackgroundTask; use binaryninja::binary_view::{BinaryView, BinaryViewExt}; use binaryninja::command::Command; use binaryninja::interaction::form::{Form, FormInputField}; use binaryninja::interaction::{MessageBoxButtonResult, MessageBoxButtonSet, MessageBoxIcon}; use binaryninja::rc::Ref; use std::path::PathBuf; use std::thread; use warp::chunk::Chunk; use warp::WarpFile; pub struct SaveFileField; impl SaveFileField { pub fn field(view: &BinaryView) -> FormInputField { let file = view.file(); let default_name = match file.project_file() { None => { // Not in a project, use the file name directly. file.display_name() } Some(project_file) => project_file.name(), }; let signature_dir = user_signature_dir(); let default_file_path = signature_dir.join(&default_name).with_extension("warp"); FormInputField::SaveFileName { prompt: "File Path".to_string(), // TODO: This is called extension but is really a filter. extension: Some("*.warp".to_string()), default: Some(default_file_path.to_string_lossy().to_string()), value: None, } } pub fn from_form(form: &Form) -> Option { let field = form.get_field_with_name("File Path")?; let field_value = field.try_value_string()?; Some(PathBuf::from(field_value)) } } pub struct OpenFileField; impl OpenFileField { pub fn field() -> FormInputField { FormInputField::OpenFileName { prompt: "Input File Path".to_string(), extension: None, default: None, value: None, } } pub fn from_form(form: &Form) -> Option { let field = form.get_field_with_name("Input File Path")?; let field_value = field.try_value_string()?; Some(PathBuf::from(field_value)) } } pub struct CreateFromCurrentView; impl CreateFromCurrentView { pub fn execute(view: Ref, external_file: bool) -> Option<()> { // Prompt the user first so that they can go do other things and not worry about a popup. let mut form = Form::new("Create From View"); if external_file { form.add_field(OpenFileField::field()); } form.add_field(SaveFileField::field(&view)); let fd_field = FileDataKindField::default(); form.add_field(fd_field.to_field()); let compression_field = CompressionTypeField::default(); form.add_field(compression_field.to_field()); let mut included_field = IncludedFunctionsField::default(); // If the view has the include tag, we better set the default to the selected functions. if view.tag_type_by_name(INCLUDE_TAG_NAME).is_some() { included_field = IncludedFunctionsField::Selected; } form.add_field(included_field.to_field()); let report_field = ReportKindField::default(); form.add_field(report_field.to_field()); let report_to_disk_field = SaveReportToDiskField::default(); form.add_field(report_to_disk_field.to_field()); if !form.prompt() { return None; } let compression_type = CompressionTypeField::from_form(&form).unwrap_or_default(); let file_path = SaveFileField::from_form(&form)?; let file_data_kind = FileDataKindField::from_form(&form).unwrap_or_default(); let file_included_functions = IncludedFunctionsField::from_form(&form).unwrap_or_default(); let report_kind = ReportKindField::from_form(&form).unwrap_or_default(); let save_report_to_disk = SaveReportToDiskField::from_form(&form).unwrap_or_default(); let open_file_path = OpenFileField::from_form(&form); // If we already have a file, prompt the user if they want to add the data. let mut existing_chunks = Vec::new(); if file_path.exists() { let prompt_result = binaryninja::interaction::show_message_box( "Keep existing file data?", "The file already exists. Do you want to keep the existing data?", MessageBoxButtonSet::YesNoCancelButtonSet, MessageBoxIcon::QuestionIcon, ); match prompt_result { MessageBoxButtonResult::NoButton => { // User wants to overwrite the file. } MessageBoxButtonResult::YesButton | MessageBoxButtonResult::OKButton => { // User wants to keep the existing data. let data = std::fs::read(&file_path).ok()?; let existing_file = WarpFile::from_owned_bytes(data)?; existing_chunks.extend(existing_file.chunks); } MessageBoxButtonResult::CancelButton => { tracing::info!( "User cancelled signature file creation, no operations were performed." ); return None; } } } let processor = WarpFileProcessor::new() .with_compression_type(compression_type) .with_file_data(file_data_kind) .with_included_functions(file_included_functions); let file = match open_file_path { None => { // We are processing the current view. NOT an external file. // Reference path is just used for the state tracking. Does not need to be readable. let reference_path = file_path.clone(); processor.process_view(reference_path, &view) } Some(open_file_path) => { // This thread will show the state in a background task. let background_task = BackgroundTask::new("Processing started...", true); new_processing_state_background_thread(background_task.clone(), processor.state()); let file = processor.process(open_file_path); background_task.finish(); file } }; if let Err(err) = file { binaryninja::interaction::show_message_box( "Failed to create signature file", &err.to_string(), MessageBoxButtonSet::OKButtonSet, MessageBoxIcon::ErrorIcon, ); tracing::error!("Failed to create signature file: {}", err); return None; } let background_task = BackgroundTask::new("Creating WARP File...", false); let mut file = file.unwrap(); // Add back the existing chunks if the user selected to keep them. if !existing_chunks.is_empty() { file.chunks.extend(existing_chunks); // TODO: Make merging optional? // TODO: Merging can lose chunk data if it goes above the maximum table count. // TODO: We should probably solve that in the warp crate itself? file.chunks = Chunk::merge(&file.chunks, compression_type.into()); // After merging, we should have at least one chunk. If not, merging actually removed data. if file.chunks.len() < 1 { tracing::error!( "Failed to merge chunks! Please report this, it should not happen." ); return None; } } let file_bytes = file.to_bytes(); let file_size = file_bytes.len(); if std::fs::write(&file_path, file_bytes).is_err() { tracing::error!("Failed to write data to signature file!"); } tracing::info!("Saved signature file to: '{}'", file_path.display()); background_task.finish(); // Show a report of the generate signatures, if desired. let report_generator = ReportGenerator::new(); if let Some(report_string) = report_generator.report(&report_kind, &file) { if save_report_to_disk == SaveReportToDiskField::Yes { let report_ext = report_generator .report_extension(&report_kind) .unwrap_or_default(); let report_path = file_path.with_extension(report_ext); let _ = std::fs::write(report_path, &report_string); } // The ReportWidget uses a QTextBrowser, which cannot render large files very well. if file_size > 10000000 { tracing::warn!("WARP report file is too large to show in the UI. Please see the report file on disk."); } else { match report_kind { ReportKindField::None => {} ReportKindField::Html => { view.show_html_report("Generated WARP File", report_string.as_str(), ""); } ReportKindField::Markdown => { view.show_markdown_report( "Generated WARP File", report_string.as_str(), "", ); } ReportKindField::Json => { view.show_plaintext_report("Generated WARP File", report_string.as_str()); } } } } Some(()) } } impl Command for CreateFromCurrentView { fn action(&self, view: &BinaryView) { let view = view.to_owned(); thread::spawn(move || { CreateFromCurrentView::execute(view, false); }); } fn valid(&self, _view: &BinaryView) -> bool { true } } pub struct CreateFromFiles; impl Command for CreateFromFiles { fn action(&self, view: &BinaryView) { let view = view.to_owned(); thread::spawn(move || { CreateFromCurrentView::execute(view, true); }); } fn valid(&self, _view: &BinaryView) -> bool { true } }