268 lines
8.4 KiB
Rust
268 lines
8.4 KiB
Rust
extern crate data_encoding;
|
|
|
|
use data_encoding::HEXUPPER;
|
|
use std::ops::Deref;
|
|
use std::str::FromStr;
|
|
|
|
pub struct BytesPattern {
|
|
bytes: Vec<u8>,
|
|
mask: Vec<bool>,
|
|
pi: Vec<usize>,
|
|
}
|
|
|
|
impl Deref for BytesPattern {
|
|
type Target = [u8];
|
|
fn deref(&self) -> &Self::Target {
|
|
&self.bytes
|
|
}
|
|
}
|
|
|
|
impl BytesPattern {
|
|
fn len(&self) -> usize {
|
|
self.deref().len()
|
|
}
|
|
fn match_at(&self, i: usize, byte: u8) -> bool {
|
|
byte == self[i] || self.mask[i]
|
|
}
|
|
fn prefix_function(&mut self) {
|
|
let mut pi = vec![0; self.len()];
|
|
let mut j = 0;
|
|
for i in 1..self.len() {
|
|
while j > 0 && !self.match_at(j, self[i]) {
|
|
j = pi[j - 1];
|
|
}
|
|
if self.match_at(j, self[i]) {
|
|
j += 1;
|
|
}
|
|
pi[i] = j;
|
|
}
|
|
self.pi = pi;
|
|
}
|
|
fn reset_offset_from(&self, current_offset: usize) -> usize {
|
|
let mut new_offset: usize = 0;
|
|
if 0 < current_offset && current_offset <= self.len() {
|
|
new_offset = self.pi[current_offset - 1];
|
|
}
|
|
new_offset
|
|
}
|
|
}
|
|
|
|
impl FromStr for BytesPattern {
|
|
type Err = &'static str;
|
|
fn from_str(str_val: &str) -> Result<Self, Self::Err> {
|
|
if str_val.is_empty() {
|
|
return Err("string is empty");
|
|
}
|
|
|
|
let mut elements: Vec<&str> = str_val.split(' ').collect();
|
|
let mask: Vec<bool> = elements.iter().map(|item| *item == "??").collect();
|
|
|
|
for (i, &item) in mask.iter().enumerate() {
|
|
if item {
|
|
elements[i] = "00";
|
|
}
|
|
}
|
|
|
|
let bytes = HEXUPPER
|
|
.decode(elements.join("").as_bytes())
|
|
.expect("decode pattern");
|
|
|
|
Ok(Self::from((bytes, mask)))
|
|
}
|
|
}
|
|
|
|
impl From<(Vec<u8>, Vec<bool>)> for BytesPattern {
|
|
fn from(x: (Vec<u8>, Vec<bool>)) -> Self {
|
|
if x.0.len() != x.1.len() {
|
|
panic!("mask length mismatch");
|
|
}
|
|
let mut s = Self {
|
|
bytes: x.0,
|
|
mask: x.1,
|
|
pi: Vec::<usize>::new(),
|
|
};
|
|
s.prefix_function();
|
|
s
|
|
}
|
|
}
|
|
|
|
pub trait FindPattern {
|
|
fn find_pattern_first(&self, pattern: &BytesPattern) -> usize;
|
|
}
|
|
|
|
pub trait ApplyPatch {
|
|
fn apply_patch(&mut self, offset: usize, patch: &[u8]);
|
|
}
|
|
|
|
impl FindPattern for &[u8] {
|
|
fn find_pattern_first(&self, pattern: &BytesPattern) -> usize {
|
|
let mut offset_result = 0;
|
|
let mut offset_pattern = 0;
|
|
let data_lenght = self.len();
|
|
let pattern_lenght = pattern.len();
|
|
|
|
for (i, &item) in self.iter().enumerate() {
|
|
if data_lenght - i < pattern_lenght - offset_pattern {
|
|
break;
|
|
}
|
|
|
|
if pattern.match_at(offset_pattern, item) {
|
|
offset_pattern += 1;
|
|
if offset_pattern == pattern_lenght {
|
|
offset_result = i + 1 - pattern_lenght;
|
|
break;
|
|
}
|
|
} else {
|
|
offset_pattern = pattern.reset_offset_from(offset_pattern);
|
|
}
|
|
}
|
|
offset_result
|
|
}
|
|
}
|
|
|
|
impl ApplyPatch for &mut [u8] {
|
|
fn apply_patch(&mut self, offset: usize, patch: &[u8]) {
|
|
for (i, &item) in patch.iter().enumerate() {
|
|
self[offset + i] = item;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
use binnpatch_macro::{binnpat, binnpats, binnvec, binnvecs};
|
|
|
|
#[test]
|
|
fn bytes_from_pattern() {
|
|
let src_pattern = "E8 4D F8 FF 83 C4 85 C0 5F 75 32 8B 54 24 48 50 50";
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
let bytes: Vec<u8> = vec![
|
|
0xE8, 0x4D, 0xF8, 0xFF, 0x83, 0xC4, 0x85, 0xC0, 0x5F, 0x75, 0x32, 0x8B, 0x54, 0x24,
|
|
0x48, 0x50, 0x50,
|
|
];
|
|
assert_eq!(bytes, pattern.bytes);
|
|
}
|
|
#[test]
|
|
fn mask_from_pattern() {
|
|
let src_pattern = "E8 ?? ?? FF ?? C4 ?? ?? 5F 75 32 ?? 54 24 48 50 50";
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
let mask: Vec<bool> = vec![
|
|
false, true, true, false, true, false, true, true, false, false, false, true, false,
|
|
false, false, false, false,
|
|
];
|
|
assert_eq!(mask, pattern.mask);
|
|
}
|
|
#[test]
|
|
fn find_pattern_first() {
|
|
let src_data: Vec<u8> = vec![
|
|
0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
|
|
];
|
|
let src_pattern = "EE FF 00 11 22";
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
let offset = src_data.as_slice().find_pattern_first(&pattern);
|
|
assert_eq!(offset, 4);
|
|
}
|
|
#[test]
|
|
fn find_pattern_first_loop() {
|
|
let src_data: Vec<u8> = vec![
|
|
0xAA, 0xBB, 0xCC, 0xAA, 0xBB, 0xEE, 0xAA, 0xBB, 0xCC, 0xAA, 0xBB, 0xCC, 0xAA, 0xBB,
|
|
0xDD, 0xAA, 0xBB, 0xCC,
|
|
];
|
|
let src_pattern = "AA BB CC AA BB DD";
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
let offset = src_data.as_slice().find_pattern_first(&pattern);
|
|
assert_eq!(offset, 9);
|
|
}
|
|
#[test]
|
|
fn prefix_function() {
|
|
let src_pattern = "AA BB BB AA AA BB BB AA BB";
|
|
let expected_pi: Vec<usize> = vec![0, 0, 0, 1, 1, 2, 3, 4, 2];
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
assert_eq!(pattern.pi, expected_pi);
|
|
}
|
|
#[test]
|
|
fn apply_patch() {
|
|
let src_data: Vec<u8> = vec![
|
|
0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
|
|
];
|
|
let expected_data: Vec<u8> = vec![
|
|
0xAA, 0xBB, 0xCC, 0xDD, 0x99, 0x88, 0x77, 0x66, 0x55, 0x33, 0x44, 0x55, 0x66,
|
|
];
|
|
let src_pattern = "EE FF 00 11 22";
|
|
let replacement = "99 88 77 66 55";
|
|
let pattern: BytesPattern = src_pattern.parse().unwrap();
|
|
let offset = src_data.as_slice().find_pattern_first(&pattern);
|
|
|
|
let replacement = replacement.replace(" ", "");
|
|
let replacement = HEXUPPER
|
|
.decode(replacement.as_bytes())
|
|
.expect("decode replacement");
|
|
let mut dst_data = src_data.clone();
|
|
dst_data.as_mut_slice().apply_patch(offset, &replacement);
|
|
assert_eq!(dst_data, expected_data);
|
|
}
|
|
|
|
#[test]
|
|
fn bytes_from_pattern_macro() {
|
|
let pattern = binnpats!("E8 4D F8 FF 83 C4 85 C0 5F 75 32 8B 54 24 48 50 50");
|
|
let bytes: Vec<u8> = vec![
|
|
0xE8, 0x4D, 0xF8, 0xFF, 0x83, 0xC4, 0x85, 0xC0, 0x5F, 0x75, 0x32, 0x8B, 0x54, 0x24,
|
|
0x48, 0x50, 0x50,
|
|
];
|
|
assert_eq!(bytes, pattern.bytes);
|
|
}
|
|
#[test]
|
|
fn mask_from_pattern_macro() {
|
|
let pattern = binnpats!("E8 ?? ?? FF ?? C4 ?? ?? 5F 75 32 ?? 54 24 48 50 50");
|
|
let mask: Vec<bool> = vec![
|
|
false, true, true, false, true, false, true, true, false, false, false, true, false,
|
|
false, false, false, false,
|
|
];
|
|
assert_eq!(mask, pattern.mask);
|
|
}
|
|
#[test]
|
|
fn binnvec_macro() {
|
|
let actual = binnvec!(E8 4D F8 FF 83 C4 85 C0 5F 75 32 8B 54 24 48 50 50);
|
|
let expected: Vec<u8> = vec![
|
|
0xE8, 0x4D, 0xF8, 0xFF, 0x83, 0xC4, 0x85, 0xC0, 0x5F, 0x75, 0x32, 0x8B, 0x54, 0x24,
|
|
0x48, 0x50, 0x50,
|
|
];
|
|
assert_eq!(expected, actual);
|
|
}
|
|
#[test]
|
|
fn binnvecs_macro() {
|
|
let actual = binnvecs!("E8 4D F8 FF 83 C4 85 C0 5F 75 32 8B 54 24 48 50 50");
|
|
let expected: Vec<u8> = vec![
|
|
0xE8, 0x4D, 0xF8, 0xFF, 0x83, 0xC4, 0x85, 0xC0, 0x5F, 0x75, 0x32, 0x8B, 0x54, 0x24,
|
|
0x48, 0x50, 0x50,
|
|
];
|
|
assert_eq!(expected, actual);
|
|
}
|
|
#[test]
|
|
fn bytes_from_pattern_macro2() {
|
|
let pattern = binnpat!(E8 4D F8 FF 83 C4 85 C0 5F 75 32 8B 54 24 48 50 50);
|
|
let bytes: Vec<u8> = vec![
|
|
0xE8, 0x4D, 0xF8, 0xFF, 0x83, 0xC4, 0x85, 0xC0, 0x5F, 0x75, 0x32, 0x8B, 0x54, 0x24,
|
|
0x48, 0x50, 0x50,
|
|
];
|
|
assert_eq!(bytes, pattern.bytes);
|
|
}
|
|
#[test]
|
|
fn mask_from_pattern_macro2() {
|
|
let pattern = binnpat!(E8 ?? ?? FF ?? C4 ?? ?? 5F 75 32 ?? 54 24 48 50 50);
|
|
let mask: Vec<bool> = vec![
|
|
false, true, true, false, true, false, true, true, false, false, false, true, false,
|
|
false, false, false, false,
|
|
];
|
|
assert_eq!(mask, pattern.mask);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "mask length mismatch")]
|
|
fn pattern_from_tupple_wrong_mask_length() {
|
|
let _pattern = BytesPattern::from((vec![1u8, 2u8, 3u8], vec![true, false]));
|
|
}
|
|
}
|