make signature type configurable

1 files changed, 328 insertions, 154 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 891c33a..a2da6b0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,32 +6,59 @@
 //! written by unknowntrojan
 //!
 
-#![feature(is_some_and, let_chains, core_intrinsics, iter_array_chunks)]
-use std::ffi::CString;
+#![feature(let_chains, core_intrinsics, iter_array_chunks)]
+use std::ffi::{CStr, CString};
+use std::fmt::Display;
 use std::ops::Range;
+use std::str::FromStr;
+use std::sync::atomic::AtomicUsize;
 use std::time::SystemTime;
 
 use binaryninja::architecture::Architecture;
 
 use binaryninja::binaryninjacore_sys::{
-    BNBinaryView, BNCreateSettings, BNFreeRelocationRanges, BNFreeSettings, BNGetRelocationRanges,
-    BNSettings, BNSettingsGetBool, BNSettingsGetUInt64, BNSettingsRegisterGroup,
-    BNSettingsRegisterSetting,
+    BNBinaryView, BNCreateSettings, BNFreeRelocationRanges, BNFreeSettings, BNFreeString,
+    BNGetRelocationRanges, BNSettings, BNSettingsGetBool, BNSettingsGetString, BNSettingsGetUInt64,
+    BNSettingsRegisterGroup, BNSettingsRegisterSetting,
 };
 use binaryninja::binaryview::{BinaryView, BinaryViewBase, BinaryViewExt};
 use binaryninja::command::{self, AddressCommand, Command};
 use clipboard::ClipboardProvider;
-use findpattern::{OwnedPattern, Pattern};
 use iced_x86::Code::{DeclareByte, DeclareDword, DeclareQword, DeclareWord};
 use iced_x86::{ConstantOffsets, FlowControl, Formatter, Instruction, NasmFormatter, OpKind};
+use rayon::prelude::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
+use rayon::slice::ParallelSlice;
+use strum::{
+    Display, EnumIter, EnumMessage, EnumString, EnumVariantNames, IntoEnumIterator, VariantNames,
+};
+
+type OwnedPattern = Vec<Option<u8>>;
+type Pattern<'a> = &'a [Option<u8>];
+
+#[derive(EnumIter, EnumVariantNames, EnumMessage, EnumString, Display)]
+enum SignatureType {
+    #[strum(message = "IDA-style signature with one ? wildcard per byte. (E9 ? ? ? ? 90)")]
+    IDAOne,
+    #[strum(message = "IDA-style signature with two ? wildcards per byte. (E9 ?? ?? ?? ?? 90)")]
+    IDATwo,
+    #[strum(message = "Rust-style signature. (0xE9, _, _, _, _, 0x90)")]
+    Rust,
+    #[strum(message = "CStr-style signature. (\"\\xE9\\x00\\x00\\x00\\x00\\x90\", \"x????x\")")]
+    CStr,
+}
+
+impl Default for SignatureType {
+    fn default() -> Self {
+        Self::IDATwo
+    }
+}
 
-struct RustSigMakerCommand;
-struct IDASigMakerCommand;
-struct CStrSigMakerCommand;
+struct SigMakerCommand;
 struct SigFinderCommand;
 
 struct RustPattern<'a>(Pattern<'a>);
-struct IDAPattern<'a>(Pattern<'a>);
+struct IDAOnePattern<'a>(Pattern<'a>);
+struct IDATwoPattern<'a>(Pattern<'a>);
 struct CStrPattern<'a>(Pattern<'a>);
 
 const MAX_INSTRUCTION_LENGTH: usize = 15;
@@ -53,7 +80,24 @@ impl<'a> core::fmt::Display for RustPattern<'a> {
     }
 }
 
-impl<'a> core::fmt::Display for IDAPattern<'a> {
+impl<'a> core::fmt::Display for IDAOnePattern<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for (i, x) in self.0.iter().enumerate() {
+            match x {
+                Some(x) => write!(f, "{:02X}", x)?,
+                None => write!(f, "?")?,
+            }
+
+            if i + 1 != self.0.len() {
+                write!(f, " ")?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl<'a> core::fmt::Display for IDATwoPattern<'a> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         for (i, x) in self.0.iter().enumerate() {
             match x {
@@ -145,23 +189,6 @@ fn get_code(bv: &BinaryView) -> Vec<(usize, Vec<u8>)> {
         .collect()
 }
 
-fn find_patterns(code_segments: &[(usize, Vec<u8>)], pattern: Pattern) -> Vec<usize> {
-    if pattern.len() == 0 {
-        return Vec::new();
-    }
-
-    code_segments
-        .into_iter()
-        .map(|segment| {
-            findpattern::find_patterns(&segment.1, &pattern)
-                .into_iter()
-                .map(|x| x + segment.0 as usize)
-                .collect::<Vec<usize>>()
-        })
-        .flatten()
-        .collect()
-}
-
 fn get_instruction_pattern(
     bv: &BinaryView,
     start_addr: usize,
@@ -259,8 +286,49 @@ fn get_instruction_pattern(
     Some(pattern)
 }
 
+fn find_patterns<'a>(
+    code_segments: &'a [(usize, Vec<u8>)],
+    pattern: Pattern<'a>,
+) -> impl ParallelIterator<Item = usize> + 'a {
+    fn find_patterns_internal_par<'a>(
+        region: &'a [u8],
+        pattern: Pattern<'a>,
+    ) -> impl ParallelIterator<Item = usize> + 'a {
+        #[inline(always)]
+        fn match_pattern(window: &[u8], pattern: Pattern) -> bool {
+            window.iter().zip(pattern).all(|(v, p)| match p {
+                Some(x) => *v == *x,
+                None => true,
+            })
+        }
+
+        region
+            .par_windows(pattern.len())
+            .enumerate()
+            .filter(|(_, wnd)| core::intrinsics::unlikely(match_pattern(wnd, pattern)))
+            .map(|(idx, _)| idx)
+    }
+
+    let pattern = pattern.clone();
+
+    code_segments
+        .par_iter()
+        .map(|segment| {
+            find_patterns_internal_par(&segment.1, pattern).map(|x| x + segment.0 as usize)
+        })
+        .flatten()
+}
+
 fn is_pattern_unique(code_segments: &[(usize, Vec<u8>)], pattern: Pattern) -> bool {
-    find_patterns(code_segments, pattern).len() == 1
+    let iter = find_patterns(code_segments, pattern);
+
+    let count = AtomicUsize::new(0);
+
+    iter.find_any(|_| {
+        count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        count.load(std::sync::atomic::Ordering::Relaxed) > 1
+    })
+    .is_none()
 }
 
 fn create_pattern_internal(
@@ -456,6 +524,31 @@ fn get_include_operands(bv: &BinaryView) -> bool {
     ret
 }
 
+fn get_signature_type(bv: &BinaryView) -> SignatureType {
+    let schema_id = CString::new("default").unwrap();
+    let key = CString::new("coolsigmaker.sig_type").unwrap();
+    let settings = unsafe { BNCreateSettings(schema_id.as_ptr()) };
+
+    let ret = unsafe {
+        BNSettingsGetString(
+            settings,
+            key.as_ptr(),
+            *std::mem::transmute::<_, *mut *mut BNBinaryView>(bv),
+            std::ptr::null_mut(),
+        )
+    };
+
+    let string = unsafe { CStr::from_ptr(ret) };
+
+    let sig_type = SignatureType::from_str(&string.to_string_lossy()).unwrap();
+
+    unsafe { BNFreeString(ret) };
+
+    unsafe { BNFreeSettings(settings) };
+
+    sig_type
+}
+
 fn register_settings() {
     fn register_setting<T>(
         settings: *mut BNSettings,
@@ -478,13 +571,50 @@ fn register_settings() {
         let properties = format!(
             r#"
 		{{
-			"title": "{}",
-			"type": "{}",
-			"default": {},
-			"description": "{}"
+			"title": "{title}",
+			"type": "{typ}",
+			"default": {default},
+			"description": "{description}",
+			"ignore": ["SettingsProjectScope", "SettingsResourceScope"]
+		}}
+		"#
+        );
+
+        let properties = CString::new(properties).unwrap();
+
+        unsafe { BNSettingsRegisterSetting(settings, name.as_ptr(), properties.as_ptr()) };
+    }
+
+    fn register_enum_setting<T>(
+        settings: *mut BNSettings,
+        name: &str,
+        title: &str,
+        description: &str,
+    ) where
+        T: Display + EnumMessage + VariantNames + IntoEnumIterator + Default,
+    {
+        let name = CString::new(name).unwrap();
+
+        let enum_variants = T::VARIANTS;
+        let enum_descriptions = T::iter()
+            .map(|x| x.get_message().unwrap_or(""))
+            .collect::<Vec<_>>();
+
+        let properties = format!(
+            r#"
+		{{
+			"title": "{title}",
+			"type": "string",
+			"default": "{}",
+			"description": "{description}",
+			"ignore": ["SettingsProjectScope", "SettingsResourceScope"],
+			"enum": {},
+			"enumDescriptions": {}
 		}}
 		"#,
-            title, typ, default, description
+            T::default(),
+            serde_json::to_string(enum_variants).unwrap(),
+            serde_json::to_string(&enum_descriptions).unwrap()
         );
 
         let properties = CString::new(properties).unwrap();
@@ -500,6 +630,7 @@ fn register_settings() {
     unsafe { BNSettingsRegisterGroup(settings, group.as_ptr(), group_fancy.as_ptr()) };
 
     register_setting::<bool>(settings, "coolsigmaker.include_operands", "Include Operands", "Include immediate operands that aren't memory-relative or relocated when creating signatures. This results in smaller, but potentially more fragile, signatures. If no unique signature can be generated without operands, we fall back to including them.", "boolean", true);
+
     register_setting::<u64>(
         settings,
         "coolsigmaker.maximum_size",
@@ -509,53 +640,167 @@ fn register_settings() {
         64,
     );
 
+    register_enum_setting::<SignatureType>(
+        settings,
+        "coolsigmaker.sig_type",
+        "Signature Type",
+        "The signature type to use for creating and finding signatures",
+    );
+
     unsafe { BNFreeSettings(settings) };
 }
 
-impl AddressCommand for RustSigMakerCommand {
-    fn action(&self, bv: &BinaryView, addr: u64) {
-        if let Some(pattern) = create_pattern(bv, addr as _) {
-            emit_result(format!("{}", RustPattern(&pattern)));
-        } else {
-            log::error!("unable to create pattern!");
-        }
-    }
+fn prepare_pattern(pattern: &str) -> String {
+    let mut pattern = pattern.to_string();
+    pattern = pattern.replace("\n", "");
+    pattern = pattern.replace("\r", "");
+    pattern = pattern.replace("\t", "");
+    pattern = pattern.trim().to_string();
 
-    fn valid(&self, bv: &BinaryView, addr: u64) -> bool {
-        is_valid(
-            bv,
-            Range {
-                start: addr,
-                end: addr + MAX_INSTRUCTION_LENGTH as u64,
-            },
-        )
-    }
+    pattern
 }
 
-impl AddressCommand for IDASigMakerCommand {
-    fn action(&self, bv: &BinaryView, addr: u64) {
-        if let Some(pattern) = create_pattern(bv, addr as _) {
-            emit_result(format!("{}", IDAPattern(&pattern)));
-        } else {
-            log::error!("unable to create pattern!");
-        }
+fn parse_idaone_pattern(pattern: &str) -> Option<OwnedPattern> {
+    // E9 ? ? ? ? 90 90
+    let mut pattern = prepare_pattern(pattern);
+
+    pattern = pattern.replace("?", "??");
+    parse_idatwo_pattern(&pattern)
+}
+
+fn parse_idatwo_pattern(pattern: &str) -> Option<OwnedPattern> {
+    // E9 ?? ?? ?? ?? 90 90
+    let mut pattern = prepare_pattern(pattern);
+    pattern = pattern.replace(" ", "");
+
+    pattern
+        .chars()
+        .array_chunks::<2>()
+        .try_fold(OwnedPattern::new(), |mut acc, byte| {
+            if byte == ['?', '?'] {
+                acc.push(None);
+            } else {
+                let Ok(byte) = u8::from_str_radix(
+				&format!("{}{}", byte[0], byte[1]),
+				16,
+			) else {
+				log::error!("unable to parse pattern!");
+				return None;
+			};
+
+                acc.push(Some(byte));
+            }
+
+            Some(acc)
+        })
+}
+
+fn parse_rust_pattern(pattern: &str) -> Option<OwnedPattern> {
+    // 0xE9, _, _, _, _, 0x90, 0x90
+    let mut pattern = prepare_pattern(pattern);
+
+    pattern = pattern.replace(",", "");
+    pattern = pattern.replace("0x", "");
+    pattern = pattern.replace("_", "??");
+
+    parse_idatwo_pattern(&pattern)
+}
+
+fn parse_cstr_pattern(pattern: &str) -> Option<OwnedPattern> {
+    // "\xE9\x00\x00\x00\x90\x90" "x????xx"
+    let pattern = prepare_pattern(pattern);
+
+    let parts = pattern.split(" ").collect::<Vec<_>>();
+
+    if parts.len() != 2 {
+        log::error!("unable to parse pattern!");
+        return None;
     }
 
-    fn valid(&self, bv: &BinaryView, addr: u64) -> bool {
-        is_valid(
-            bv,
-            Range {
-                start: addr,
-                end: addr + MAX_INSTRUCTION_LENGTH as u64,
-            },
-        )
+    let first_part = parts[0].trim().replace("\"", "");
+
+    let first_part = first_part
+        .split("\\x")
+        .into_iter()
+        .map(|x| {
+            if x == "" || x.len() < 2 {
+                vec![]
+            } else {
+                let byte = &x[..2];
+                let wildcards = x.len() - 2;
+
+                let wildcards = (0..wildcards)
+                    .into_iter()
+                    .map(|_| None)
+                    .collect::<Vec<Option<u8>>>();
+
+                let mut result = vec![u8::from_str_radix(&byte, 16)
+                    .map_err(|x| {
+                        log::error!("unable to parse pattern!");
+                        x
+                    })
+                    .ok()];
+
+                result.extend(wildcards);
+
+                result
+            }
+        })
+        .flatten()
+        .collect::<Vec<_>>();
+
+    let second_part = parts[1].trim().replace("\"", "");
+
+    let mut pattern = OwnedPattern::new();
+
+    for (byte, mask) in first_part.iter().zip(second_part.chars()) {
+        match mask {
+            'x' => {
+                pattern.push(Some(byte.unwrap()));
+            }
+            '?' => {
+                pattern.push(None);
+            }
+            _ => {
+                log::error!("invalid char encountered!");
+            }
+        }
     }
+
+    Some(pattern)
+}
+
+#[test]
+fn test_patterns() {
+    let test_pattern = &[Some(0xE9), None, None, None, None, Some(0x90), Some(0x90)];
+    assert_eq!(
+        &parse_idaone_pattern("E9 ? ? ? ? 90 90").unwrap(),
+        test_pattern
+    );
+    assert_eq!(
+        &parse_idatwo_pattern("E9 ?? ?? ?? ?? 90 90").unwrap(),
+        test_pattern
+    );
+    assert_eq!(
+        &parse_rust_pattern("0xE9, _, _, _, _, 0x90, 0x90").unwrap(),
+        test_pattern
+    );
+    assert_eq!(
+        &parse_cstr_pattern(r#""\xE9\x00??\x00\x90\x90" "x????xx""#).unwrap(),
+        test_pattern
+    );
 }
 
-impl AddressCommand for CStrSigMakerCommand {
+impl AddressCommand for SigMakerCommand {
     fn action(&self, bv: &BinaryView, addr: u64) {
         if let Some(pattern) = create_pattern(bv, addr as _) {
-            emit_result(format!("{}", CStrPattern(&pattern)));
+            let pattern: Box<dyn core::fmt::Display> = match get_signature_type(bv) {
+                SignatureType::IDAOne => Box::new(IDAOnePattern(&pattern)),
+                SignatureType::IDATwo => Box::new(IDATwoPattern(&pattern)),
+                SignatureType::Rust => Box::new(RustPattern(&pattern)),
+                SignatureType::CStr => Box::new(CStrPattern(&pattern)),
+            };
+            emit_result(format!("{}", pattern));
         } else {
             log::error!("unable to create pattern!");
         }
@@ -574,84 +819,25 @@ impl AddressCommand for CStrSigMakerCommand {
 
 impl Command for SigFinderCommand {
     fn action(&self, bv: &BinaryView) {
-        // please don't look a this code. its absolutely disgusting. i hate working with strings so much.
-
-        let Ok(mut sig) = get_clipboard_contents() else {
+        let Ok(sig) = get_clipboard_contents() else {
 			log::error!("unable to get signature from clipboard!");
 			return;
 		};
 
         let data = get_code(bv);
 
-        sig = sig.replace("\n", "");
-        sig = sig.replace("\r", "");
-        sig = sig.replace("\t", "");
-        sig = sig.trim().to_string();
-
-        let mut pattern = OwnedPattern::new();
-
-        if sig.contains("\"") {
-            let parts = sig.split(" ").collect::<Vec<_>>();
-
-            if parts.len() != 2 {
-                log::error!("unable to parse pattern!");
-                return;
-            }
-
-            let first_part = parts[0].trim().replace("\"", "");
-
-            let first_part = first_part
-                .split("\\x")
-                .into_iter()
-                .filter_map(|x| if x == "" { None } else { Some(x.split("?")) })
-                .flatten()
-                .collect::<Vec<_>>();
-
-            let second_part = parts[1].trim().replace("\"", "");
-
-            for (byte, mask) in first_part.iter().zip(second_part.chars()) {
-                match mask {
-                    'x' => {
-                        let Ok(byte) = u8::from_str_radix(byte, 16) else {
-							log::error!("unable to parse pattern!");
-							return;
-						};
-                        pattern.push(Some(byte));
-                    }
-                    '?' => {
-                        pattern.push(None);
-                    }
-                    _ => {
-                        log::error!("invalid char encountered!");
-                    }
-                }
-            }
-        } else {
-            sig = sig.replace("_", "??");
-            sig = sig.replace("0x", "");
-            sig = sig.replace(", ", "");
-            sig = sig.replace(" ", "");
-
-            for byte in sig.chars().array_chunks::<2>() {
-                if byte == ['?', '?'] {
-                    pattern.push(None);
-                } else {
-                    let Ok(byte) = u8::from_str_radix(
-						&format!("{}{}", byte[0], byte[1]),
-						16,
-					) else {
-						log::error!("unable to parse pattern!");
-						return;
-					};
-
-                    pattern.push(Some(byte));
-                }
-            }
-        }
+        let Some(pattern) = (match get_signature_type(bv) {
+            SignatureType::IDAOne => parse_idaone_pattern(&sig),
+            SignatureType::IDATwo => parse_idatwo_pattern(&sig),
+            SignatureType::Rust => parse_rust_pattern(&sig),
+            SignatureType::CStr => parse_cstr_pattern(&sig),
+        }) else {
+			log::error!("failed to parse pattern.");
+			return;
+		};
 
-        for occurrence in find_patterns(&data, &pattern) {
-            log::info!("found signature at {:#04X}", occurrence);
-        }
+        find_patterns(&data, &pattern)
+            .for_each(|occurrence| log::info!("found signature at {:#04X}", occurrence));
 
         log::info!("scan finished.");
     }
@@ -686,21 +872,9 @@ pub extern "C" fn CorePluginInit() -> bool {
     register_settings();
 
     command::register_for_address(
-        "CoolSigMaker\\Create Signature from Address: Rust (0xE9, _, _, _, _)",
-        "Creates a Rust-style signature from the currently selected address",
-        RustSigMakerCommand {},
-    );
-
-    command::register_for_address(
-        "CoolSigMaker\\Create Signature from Address: IDA (E9 ?? ?? ?? ??)",
-        "Creates an IDA-style signature from the currently selected address",
-        IDASigMakerCommand {},
-    );
-
-    command::register_for_address(
-        "CoolSigMaker\\Create Signature from Address: CStr (\"/xE9/x00/x00/x00/x00\" \"x????\")",
-        "Creates a CStr-style signature from the currently selected address",
-        CStrSigMakerCommand {},
+        "CoolSigMaker\\Create Signature from Address",
+        "Creates a Signature from the currently selected address",
+        SigMakerCommand {},
     );
 
     command::register(