use std::env;
use rand::{Rng, thread_rng};
fn main() {
let mut args = env::args_os();
args.next();
if args.len() <= 1 {
println!("markovbin <COUNT> <FILES>...\n\nGiven a list of binary <FILES> (.exe, .dll) analyzes their code to generate a markov chain of <COUNT> bytes.");
return;
}
let count: usize = match args.next().unwrap().into_string().unwrap().parse() {
Ok(count) => count,
Err(e) => {
eprintln!("first argument should be a number (of bytes to generate): {}", e);
return;
}
};
let mut markov_data = [[0u32; 256]; 256];
for file in args {
eprintln!("generating markov data from {:?}...", file);
match pelite::FileMap::open(&file) {
Ok(map) => {
match pelite::PeFile::from_bytes(&map) {
Ok(pe) => {
for section in pe.section_headers() {
if section.Characteristics & 0xE0000000 == 0x60000000 {
match pe.derva_slice(section.VirtualAddress, section.VirtualSize as usize) {
Ok(bytes) => analyze(bytes, &mut markov_data),
Err(e) => eprintln!("cannot read code from {:?}: {}", file, e),
}
}
}
},
Err(err) => {
eprintln!("cannot parse file {:?}: {}", file, err);
},
}
},
Err(e) => {
eprintln!("cannot read file {:?}: {}", file, e);
},
}
}
generate(&markov_data, count);
}
fn analyze(bytes: &[u8], buckets: &mut [[u32; 256]; 256]) {
if bytes.len() == 0 {
return;
}
let mut bucket = bytes[0] as usize;
for &byte in &bytes[1..] {
buckets[bucket][byte as usize] += 1;
bucket = byte as usize;
}
}
fn generate(buckets: &[[u32; 256]; 256], count: usize) {
let mut rng = thread_rng();
let mut byte = rng.gen();
for _ in 0..count {
let bucket = &buckets[byte as usize];
let total = bucket.iter().sum();
if total == 0 {
byte = rng.gen();
continue;
}
print!("{:02X} ", byte);
let mut pick = rng.gen_range(0, total);
for i in 0..256 {
if pick < bucket[i] {
byte = i as u8;
break;
}
pick -= bucket[i];
}
}
println!();
}