balrog

~ ELF file analysis tool to be used after firmware extraction
git clone https://github.com/djnnvx/balrog.git
mod.rs (7907B)
      1 pub mod aarch64;
      2 pub mod arm;
      3 pub mod mips;
      4 
      5 use crate::load::Binary;
      6 use anyhow::Result;
      7 use capstone::arch::arm::ArmOperandType;
      8 use capstone::arch::arm64::Arm64OperandType;
      9 use capstone::arch::mips::MipsOperand;
     10 use capstone::arch::ArchOperand;
     11 use capstone::prelude::*;
     12 use capstone::{Endian, Insn as CsInsn};
     13 use std::collections::HashMap;
     14 
     15 const IMM16_MASK: u64 = 0xffff;
     16 const HIGH_HALF_SHIFT: u32 = 16;
     17 
     18 #[derive(Clone, Copy, PartialEq, Eq, Debug, serde::Serialize)]
     19 pub enum Arch {
     20     Arm,
     21     Aarch64,
     22     Mips,
     23     X86,
     24     X86_64,
     25 }
     26 
     27 pub const ANALYZABLE_ARCHES: &str = "mips, arm, aarch64";
     28 
     29 impl Arch {
     30     pub fn align(self) -> u64 {
     31         match self {
     32             Arch::X86 | Arch::X86_64 => 1,
     33             _ => 4,
     34         }
     35     }
     36     pub fn name(self) -> &'static str {
     37         match self {
     38             Arch::Arm => "arm",
     39             Arch::Aarch64 => "aarch64",
     40             Arch::Mips => "mips",
     41             Arch::X86 => "x86",
     42             Arch::X86_64 => "x86_64",
     43         }
     44     }
     45     pub fn analyzable(self) -> bool {
     46         matches!(self, Arch::Arm | Arch::Aarch64 | Arch::Mips)
     47     }
     48 }
     49 
     50 #[derive(Clone, Copy, Debug)]
     51 pub enum Kind {
     52     Call(u64),
     53     CallInd { reg: u16 },
     54     SetConst { dst: u16, val: u64 },
     55     Load { dst: u16, base: u16, disp: i64 },
     56     AddImm { dst: u16, base: u16, imm: i64 },
     57     PcVal { dst: u16, val: u64 },
     58     Mov { dst: u16, src: u16 },
     59     Ret,
     60     Other,
     61 }
     62 
     63 pub struct Ins {
     64     pub addr: u64,
     65     pub kind: Kind,
     66 }
     67 
     68 pub fn build(arch: Arch, little: bool) -> Result<Capstone> {
     69     let endian = if little { Endian::Little } else { Endian::Big };
     70     let cs = match arch {
     71         Arch::Arm => Capstone::new()
     72             .arm()
     73             .mode(capstone::arch::arm::ArchMode::Arm)
     74             .endian(endian)
     75             .detail(true)
     76             .build(),
     77         Arch::Aarch64 => Capstone::new()
     78             .arm64()
     79             .mode(capstone::arch::arm64::ArchMode::Arm)
     80             .endian(endian)
     81             .detail(true)
     82             .build(),
     83         Arch::Mips => Capstone::new()
     84             .mips()
     85             .mode(capstone::arch::mips::ArchMode::Mips32)
     86             .endian(endian)
     87             .detail(true)
     88             .build(),
     89         Arch::X86 => Capstone::new()
     90             .x86()
     91             .mode(capstone::arch::x86::ArchMode::Mode32)
     92             .detail(true)
     93             .build(),
     94         Arch::X86_64 => Capstone::new()
     95             .x86()
     96             .mode(capstone::arch::x86::ArchMode::Mode64)
     97             .detail(true)
     98             .build(),
     99     }?;
    100     Ok(cs)
    101 }
    102 
    103 pub fn sweep(cs: &Capstone, bin: &Binary) -> Vec<Ins> {
    104     let align = bin.arch.align();
    105     let mut out = Vec::new();
    106     for seg in bin.segs.iter().filter(|seg| seg.exec) {
    107         let start = seg.off as usize;
    108         let end = ((seg.off + seg.filesz) as usize).min(bin.data.len());
    109         if start >= end {
    110             continue;
    111         }
    112         let code = &bin.data[start..end];
    113         let mut pos: u64 = 0;
    114         let len = (end - start) as u64;
    115         while pos < len {
    116             let Ok(insns) = cs.disasm_all(&code[pos as usize..], seg.vaddr + pos) else {
    117                 pos += align;
    118                 continue;
    119             };
    120             if insns.is_empty() {
    121                 pos += align;
    122                 continue;
    123             }
    124             let mut last_end = pos;
    125             for insn in insns.iter() {
    126                 out.push(lower(cs, bin, insn));
    127                 last_end = (insn.address() - seg.vaddr) + insn.len() as u64;
    128             }
    129             pos = last_end.max(pos + align);
    130         }
    131     }
    132     out.sort_by_key(|ins| ins.addr);
    133     out.dedup_by_key(|ins| ins.addr);
    134     out
    135 }
    136 
    137 fn lower(cs: &Capstone, bin: &Binary, insn: &CsInsn) -> Ins {
    138     Ins {
    139         addr: insn.address(),
    140         kind: classify(cs, bin, insn).unwrap_or(Kind::Other),
    141     }
    142 }
    143 
    144 #[derive(Clone, Copy)]
    145 pub enum Op {
    146     Reg(u16),
    147     Imm(i64),
    148     Mem(u16, i64),
    149     None,
    150 }
    151 
    152 fn operands(ops: &[ArchOperand]) -> [Op; 3] {
    153     std::array::from_fn(|n| match ops.get(n) {
    154         Some(ArchOperand::MipsOperand(MipsOperand::Reg(r))) => Op::Reg(r.0),
    155         Some(ArchOperand::MipsOperand(MipsOperand::Imm(v))) => Op::Imm(*v),
    156         Some(ArchOperand::MipsOperand(MipsOperand::Mem(m))) => Op::Mem(m.base().0, m.disp()),
    157         Some(ArchOperand::ArmOperand(o)) => match o.op_type {
    158             ArmOperandType::Reg(r) => Op::Reg(r.0),
    159             ArmOperandType::Imm(v) => Op::Imm(v as i64),
    160             ArmOperandType::Mem(m) => Op::Mem(m.base().0, m.disp() as i64),
    161             _ => Op::None,
    162         },
    163         Some(ArchOperand::Arm64Operand(o)) => match o.op_type {
    164             Arm64OperandType::Reg(r) => Op::Reg(r.0),
    165             Arm64OperandType::Imm(v) => Op::Imm(v),
    166             Arm64OperandType::Mem(m) => Op::Mem(m.base().0, m.disp() as i64),
    167             _ => Op::None,
    168         },
    169         _ => Op::None,
    170     })
    171 }
    172 
    173 fn reg(o: &[Op; 3], n: usize) -> Option<u16> {
    174     match o[n] {
    175         Op::Reg(r) => Some(r),
    176         _ => None,
    177     }
    178 }
    179 
    180 fn imm(o: &[Op; 3], n: usize) -> Option<i64> {
    181     match o[n] {
    182         Op::Imm(v) => Some(v),
    183         _ => None,
    184     }
    185 }
    186 
    187 fn mem(o: &[Op; 3], n: usize) -> Option<(u16, i64)> {
    188     match o[n] {
    189         Op::Mem(b, d) => Some((b, d)),
    190         _ => None,
    191     }
    192 }
    193 
    194 fn classify(cs: &Capstone, bin: &Binary, insn: &CsInsn) -> Option<Kind> {
    195     let detail = cs.insn_detail(insn).ok()?;
    196     let ops = operands(&detail.arch_detail().operands());
    197     let mnemonic = insn.mnemonic()?;
    198     match bin.arch {
    199         Arch::Mips => mips::decode(mnemonic, &ops),
    200         Arch::Aarch64 => aarch64::decode(mnemonic, &ops),
    201         Arch::Arm => arm::decode(mnemonic, &ops, insn.address(), bin),
    202         _ => None,
    203     }
    204 }
    205 
    206 fn arg_regs(arch: Arch) -> &'static [u16] {
    207     match arch {
    208         Arch::Mips => &mips::ARG_REGS,
    209         Arch::Aarch64 => &aarch64::ARG_REGS,
    210         Arch::Arm => &arm::ARG_REGS,
    211         _ => &[],
    212     }
    213 }
    214 
    215 pub fn arg_reg(arch: Arch, n: usize) -> Option<u16> {
    216     arg_regs(arch).get(n).copied()
    217 }
    218 
    219 pub fn arg_count(arch: Arch) -> usize {
    220     arg_regs(arch).len()
    221 }
    222 
    223 pub fn gp_reg(arch: Arch) -> Option<u16> {
    224     match arch {
    225         Arch::Mips => Some(mips::GP),
    226         _ => None,
    227     }
    228 }
    229 
    230 pub fn sp_reg(arch: Arch) -> u16 {
    231     match arch {
    232         Arch::Mips => mips::SP,
    233         Arch::Aarch64 => aarch64::SP,
    234         Arch::Arm => arm::SP,
    235         _ => 0,
    236     }
    237 }
    238 
    239 pub fn resolve_stubs(bin: &Binary, insns: &[Ins]) -> HashMap<u64, String> {
    240     let mut out = HashMap::new();
    241     for (name, addr) in &bin.undef_stub {
    242         out.insert(*addr, name.clone());
    243     }
    244     let mut tracked: HashMap<u16, (u64, u64)> = HashMap::new();
    245     for ins in insns {
    246         match ins.kind {
    247             Kind::SetConst { dst, val } => {
    248                 tracked.insert(dst, (ins.addr, val));
    249             }
    250             Kind::AddImm { dst, base, imm } => {
    251                 if let Some(&(origin, value)) = tracked.get(&base) {
    252                     tracked.insert(dst, (origin, value.wrapping_add(imm as u64)));
    253                 } else {
    254                     tracked.remove(&dst);
    255                 }
    256             }
    257             Kind::Load { dst, base, disp } => {
    258                 if let Some(&(origin, value)) = tracked.get(&base) {
    259                     let slot = value.wrapping_add(disp as u64);
    260                     if let Some(name) = bin.got2sym.get(&slot) {
    261                         out.insert(origin, name.clone());
    262                     }
    263                 }
    264                 tracked.remove(&dst);
    265             }
    266             Kind::PcVal { dst, val } => {
    267                 tracked.insert(dst, (ins.addr, val));
    268             }
    269             Kind::Mov { dst, src } => {
    270                 if let Some(&origin_value) = tracked.get(&src) {
    271                     tracked.insert(dst, origin_value);
    272                 } else {
    273                     tracked.remove(&dst);
    274                 }
    275             }
    276             _ => {}
    277         }
    278     }
    279     out
    280 }