patch_bin.rs

extern crate data_encoding;
extern crate memmap;

use data_encoding::HEXUPPER;
use memmap::{Mmap, MmapMut};
use std::fs;
use std::fs::{File, OpenOptions};
use std::ops::{Deref, DerefMut};

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 From<&str> for BytesPattern {
    fn from(str_val: &str) -> Self {
        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 mut s = Self {
            bytes: HEXUPPER
                .decode(elements.join("").as_bytes())
                .expect("decode pattern"),
            mask,
            pi: Vec::<usize>::new(),
        };
        s.prefix_function();
        s
    }
}

trait FindPattern {
    fn find_pattern_first(&self, pattern: &BytesPattern) -> usize;
}

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;
        }
    }
}

fn main() {
    let src_file = "Some.dll";
    let dst_file = "Some.dll.patched";
    let src_pattern =
        "40 3E 1D ?? ?? 12 1C 7C 48 ?? 73 6F 02 22 ?? 61 19 4E 13 60 48 45 19 27 5B";
    let replacement =
        "06 5A 18 74 2D 62 12 6A 13 4A 2B 0E 6F 0F 36 7A 28 0A 37 67 0A 4B 01 73 14";

    let pattern = BytesPattern::from(src_pattern);

    let replacement = replacement.replace(" ", "");
    let replacement = HEXUPPER
        .decode(replacement.as_bytes())
        .expect("decode replacement");

    let file = File::open(src_file).expect("src open");

    let src_map = unsafe { Mmap::map(&file) }.expect("src map");
    let data_src = src_map.deref();

    let offset = data_src.find_pattern_first(&pattern);
    if 0 == offset {
        panic!("src pattern not found");
    }
    println!("found at {:#010X}", offset);

    if replacement.len() > data_src.len() - offset {
        panic!("replace overflow");
    }

    drop(src_map);
    drop(file);

    fs::copy(src_file, dst_file).expect("file copy");

    let file = OpenOptions::new()
        .read(true)
        .write(true)
        .open(dst_file)
        .expect("dst open");
    let mut dst_map = unsafe { MmapMut::map_mut(&file) }.expect("dst map");
    let mut data_dst = dst_map.deref_mut();

    data_dst.apply_patch(offset, &replacement);
    drop(dst_map);

    println!("done");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[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::from(src_pattern);
        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::from(src_pattern);
        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::from(src_pattern);
        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::from(src_pattern);
        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::from(src_pattern);
        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::from(src_pattern);
        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);
    }
}