balrog

~ ELF file analysis tool to be used after firmware extraction
git clone https://github.com/djnnvx/balrog.git
graph.rs (9252B)
      1 use crate::arch::{self, Ins, Kind};
      2 use crate::load::Binary;
      3 use std::collections::hash_map::Entry;
      4 use std::collections::{HashMap, HashSet, VecDeque};
      5 
      6 #[derive(Clone, Copy, Debug, PartialEq)]
      7 pub enum Val {
      8     Num(u64),
      9     Stack(i64),
     10     Import(u64),
     11 }
     12 
     13 pub struct CallSite {
     14     pub site: u64,
     15     pub func: u64,
     16     pub import: Option<String>,
     17     pub target: u64,
     18     pub args: Vec<Option<Val>>,
     19 }
     20 
     21 pub struct Graph {
     22     pub names: HashMap<u64, String>,
     23     pub sites: Vec<CallSite>,
     24     pub edges: HashMap<u64, HashSet<u64>>,
     25     pub strings: HashMap<u64, Vec<String>>,
     26     pub addr_taken: HashSet<u64>,
     27 }
     28 
     29 fn function_starts(bin: &Binary, insns: &[Ins], stubs: &HashMap<u64, String>) -> Vec<u64> {
     30     let decoded: HashSet<u64> = insns.iter().map(|ins| ins.addr).collect();
     31     let mut starts: HashSet<u64> = HashSet::new();
     32     if decoded.contains(&bin.entry) {
     33         starts.insert(bin.entry);
     34     }
     35     for (_, addr) in bin.locals.iter().chain(bin.exports.iter()) {
     36         if decoded.contains(addr) {
     37             starts.insert(*addr);
     38         }
     39     }
     40     for ins in insns {
     41         if let Kind::Call(target) = ins.kind {
     42             if !stubs.contains_key(&target) && decoded.contains(&target) {
     43                 starts.insert(target);
     44             }
     45         }
     46     }
     47     for addr in stubs.keys() {
     48         starts.remove(addr);
     49     }
     50     let mut starts: Vec<u64> = starts.into_iter().collect();
     51     starts.sort_unstable();
     52     starts
     53 }
     54 
     55 fn function_names(bin: &Binary) -> HashMap<u64, String> {
     56     let mut names = HashMap::new();
     57     for (name, addr) in bin.locals.iter().chain(bin.exports.iter()) {
     58         names.entry(*addr).or_insert_with(|| name.clone());
     59     }
     60     names
     61 }
     62 
     63 struct Regs<'a> {
     64     bin: &'a Binary,
     65     sp: u16,
     66     gp: Option<u16>,
     67     nargs: usize,
     68     vals: HashMap<u16, Val>,
     69 }
     70 
     71 impl<'a> Regs<'a> {
     72     fn new(bin: &'a Binary) -> Regs<'a> {
     73         Regs {
     74             bin,
     75             sp: arch::sp_reg(bin.arch),
     76             gp: arch::gp_reg(bin.arch),
     77             nargs: arch::arg_count(bin.arch),
     78             vals: HashMap::new(),
     79         }
     80     }
     81 
     82     fn enter(&mut self) {
     83         self.vals.clear();
     84         self.vals.insert(self.sp, Val::Stack(0));
     85         self.pin_gp();
     86     }
     87 
     88     fn pin_gp(&mut self) {
     89         if let (Some(g), Some(r)) = (self.bin.gp, self.gp) {
     90             self.vals.insert(r, Val::Num(g));
     91         }
     92     }
     93 
     94     fn set(&mut self, dst: u16, v: Option<Val>) -> Option<Val> {
     95         match v {
     96             Some(v) => self.vals.insert(dst, v),
     97             None => self.vals.remove(&dst),
     98         };
     99         v
    100     }
    101 
    102     fn got_slot(&self, base: u16, disp: i64) -> Option<u64> {
    103         match self.vals.get(&base) {
    104             Some(Val::Num(addr)) => Some(addr.wrapping_add(disp as u64)),
    105             _ => None,
    106         }
    107     }
    108 
    109     fn step(&mut self, kind: Kind) -> Option<Val> {
    110         let wrote = match kind {
    111             Kind::SetConst { dst, val } | Kind::PcVal { dst, val } => {
    112                 self.set(dst, Some(Val::Num(val)))
    113             }
    114             Kind::AddImm { dst, base, imm } => {
    115                 let sum = match self.vals.get(&base) {
    116                     Some(Val::Num(addr)) => Some(Val::Num(addr.wrapping_add(imm as u64))),
    117                     Some(Val::Stack(offset)) => Some(Val::Stack(offset + imm)),
    118                     _ => None,
    119                 };
    120                 self.set(dst, sum)
    121             }
    122             Kind::Mov { dst, src } => {
    123                 let moved = self.vals.get(&src).copied();
    124                 self.set(dst, moved)
    125             }
    126             Kind::Load { dst, base, disp } => {
    127                 let import = self
    128                     .got_slot(base, disp)
    129                     .filter(|slot| self.bin.got2sym.contains_key(slot))
    130                     .map(Val::Import);
    131                 self.set(dst, import)
    132             }
    133             _ => None,
    134         };
    135         self.pin_gp();
    136         wrote
    137     }
    138 
    139     fn arg_value(&self, index: usize) -> Option<Val> {
    140         let reg = arch::arg_reg(self.bin.arch, index)?;
    141         let value = self.vals.get(&reg).copied()?;
    142         if matches!(value, Val::Import(_)) {
    143             return None;
    144         }
    145         Some(value)
    146     }
    147 
    148     fn args(&self) -> Vec<Option<Val>> {
    149         (0..self.nargs).map(|index| self.arg_value(index)).collect()
    150     }
    151 
    152     fn import_at(&self, reg: u16) -> Option<&String> {
    153         match self.vals.get(&reg) {
    154             Some(Val::Import(slot)) => self.bin.got2sym.get(slot),
    155             _ => None,
    156         }
    157     }
    158 
    159     fn clobber_args(&mut self) {
    160         for n in 0..self.nargs {
    161             if let Some(r) = arch::arg_reg(self.bin.arch, n) {
    162                 self.vals.remove(&r);
    163             }
    164         }
    165     }
    166 }
    167 
    168 const MAX_STRINGS_PER_FUNC: usize = 24;
    169 const MIN_STRING_LEN: usize = 4;
    170 
    171 fn record_constant(bin: &Binary, func: u64, value: u64, starts: &[u64], out: &mut Graph) {
    172     if starts.binary_search(&value).is_ok() {
    173         out.addr_taken.insert(value);
    174         return;
    175     }
    176     let Some(text) = bin.cstr_at(value) else {
    177         return;
    178     };
    179     if text.len() < MIN_STRING_LEN {
    180         return;
    181     }
    182     let strings = out.strings.entry(func).or_default();
    183     if strings.len() < MAX_STRINGS_PER_FUNC && !strings.contains(&text) {
    184         strings.push(text);
    185     }
    186 }
    187 
    188 fn first_arg_is_constant(site: &CallSite) -> bool {
    189     matches!(site.args.first(), Some(Some(Val::Num(_))))
    190 }
    191 
    192 impl Graph {
    193     pub fn build(bin: &Binary, insns: &[Ins], stubs: HashMap<u64, String>) -> Graph {
    194         let starts = function_starts(bin, insns, &stubs);
    195         let mut g = Graph {
    196             names: function_names(bin),
    197             sites: Vec::new(),
    198             edges: HashMap::new(),
    199             strings: HashMap::new(),
    200             addr_taken: HashSet::new(),
    201         };
    202         let mut regs = Regs::new(bin);
    203         let mut cur_func: u64 = starts.first().copied().unwrap_or(0);
    204 
    205         for ins in insns {
    206             if starts.binary_search(&ins.addr).is_ok() {
    207                 cur_func = ins.addr;
    208                 regs.enter();
    209             }
    210             let wrote = regs.step(ins.kind);
    211             match ins.kind {
    212                 Kind::Call(target) => {
    213                     let import = stubs.get(&target).cloned();
    214                     if import.is_none() && starts.binary_search(&target).is_ok() {
    215                         g.edges.entry(cur_func).or_default().insert(target);
    216                     }
    217                     g.sites.push(CallSite {
    218                         site: ins.addr,
    219                         func: cur_func,
    220                         import,
    221                         target,
    222                         args: regs.args(),
    223                     });
    224                     regs.clobber_args();
    225                 }
    226                 Kind::CallInd { reg } => {
    227                     if let Some(name) = regs.import_at(reg).cloned() {
    228                         g.sites.push(CallSite {
    229                             site: ins.addr,
    230                             func: cur_func,
    231                             import: Some(name),
    232                             target: 0,
    233                             args: regs.args(),
    234                         });
    235                     }
    236                     regs.clobber_args();
    237                 }
    238                 _ => {}
    239             }
    240             if let Some(Val::Num(value)) = wrote {
    241                 record_constant(bin, cur_func, value, &starts, &mut g);
    242             }
    243         }
    244         g
    245     }
    246 
    247     pub fn func_name(&self, addr: u64) -> String {
    248         self.names
    249             .get(&addr)
    250             .cloned()
    251             .unwrap_or_else(|| format!("sub_{addr:x}"))
    252     }
    253 
    254     pub fn reach_from(&self, inputs: &[&str]) -> HashMap<u64, u32> {
    255         let mut dist: HashMap<u64, u32> = HashMap::new();
    256         let mut queue = VecDeque::new();
    257         for site in &self.sites {
    258             let Some(import) = &site.import else { continue };
    259             if inputs.contains(&import.as_str()) {
    260                 dist.insert(site.func, 0);
    261                 queue.push_back(site.func);
    262             }
    263         }
    264         while let Some(func) = queue.pop_front() {
    265             let depth = dist[&func];
    266             let Some(callees) = self.edges.get(&func) else {
    267                 continue;
    268             };
    269             for &callee in callees {
    270                 if let Entry::Vacant(slot) = dist.entry(callee) {
    271                     slot.insert(depth + 1);
    272                     queue.push_back(callee);
    273                 }
    274             }
    275         }
    276         dist
    277     }
    278 
    279     pub fn shell_equivalent(&self, seeds: &[&str]) -> HashSet<u64> {
    280         let mut tainted: HashSet<u64> = HashSet::new();
    281         for site in &self.sites {
    282             let Some(import) = &site.import else { continue };
    283             if seeds.contains(&import.as_str()) && !first_arg_is_constant(site) {
    284                 tainted.insert(site.func);
    285             }
    286         }
    287         loop {
    288             let before = tainted.len();
    289             for site in &self.sites {
    290                 if site.import.is_none()
    291                     && tainted.contains(&site.target)
    292                     && !first_arg_is_constant(site)
    293                 {
    294                     tainted.insert(site.func);
    295                 }
    296             }
    297             if tainted.len() == before {
    298                 break;
    299             }
    300         }
    301         tainted
    302     }
    303 }