balrog

~ ELF file analysis tool to be used after firmware extraction
git clone https://github.com/djnnvx/balrog.git
main.rs (8913B)
      1 mod arch;
      2 mod graph;
      3 mod load;
      4 mod rules;
      5 
      6 use anyhow::Result;
      7 use clap::{Args, Parser, Subcommand};
      8 use rayon::prelude::*;
      9 use std::cmp::Ordering;
     10 use std::collections::{HashMap, HashSet};
     11 use std::io::{self, BufWriter, Write};
     12 use std::path::{Path, PathBuf};
     13 
     14 const ELF_MAGIC: [u8; 4] = [0x7f, b'E', b'L', b'F'];
     15 const STUB_PREVIEW: usize = 12;
     16 
     17 #[derive(Parser)]
     18 #[command(
     19     name = "balrog",
     20     version,
     21     about = "risky-call and reachability triage for firmware ELF trees"
     22 )]
     23 struct Cli {
     24     #[command(subcommand)]
     25     cmd: Cmd,
     26 }
     27 
     28 #[derive(Subcommand)]
     29 enum Cmd {
     30     #[command(about = "rank risky calls in a binary or a whole extracted tree")]
     31     Scan(ScanArgs),
     32     #[command(about = "resolve imports and their PLT/lazy-binding stubs")]
     33     Imports { path: PathBuf },
     34 }
     35 
     36 #[derive(Args)]
     37 struct ScanArgs {
     38     #[arg(
     39         required = true,
     40         help = "one or more binaries or trees. Firmware has tiers: rootfs and app bundles are separate trees and both need triaging"
     41     )]
     42     paths: Vec<PathBuf>,
     43     #[arg(short, long)]
     44     json: bool,
     45     #[arg(short = 'n', long = "top", default_value_t = 15)]
     46     per_class_limit: usize,
     47     #[arg(long = "min", default_value_t = 40, help = "minimum score to report")]
     48     min_score: i32,
     49     #[arg(
     50         long = "per-bin",
     51         default_value_t = 2,
     52         help = "at most this many findings per class per binary"
     53     )]
     54     max_per_binary: usize,
     55     #[arg(long = "class", value_enum, help = "only report this class")]
     56     only_class: Option<rules::Class>,
     57 }
     58 
     59 fn is_elf(path: &Path) -> bool {
     60     let Ok(mut file) = std::fs::File::open(path) else {
     61         return false;
     62     };
     63     use std::io::Read;
     64     let mut magic = [0u8; 4];
     65     file.read_exact(&mut magic).is_ok() && magic == ELF_MAGIC
     66 }
     67 
     68 fn collect_elf_files(root: &Path) -> Vec<PathBuf> {
     69     if root.is_file() {
     70         return vec![root.to_path_buf()];
     71     }
     72     let mut elf_files = Vec::new();
     73     let mut unvisited_dirs = vec![root.to_path_buf()];
     74     while let Some(dir) = unvisited_dirs.pop() {
     75         let Ok(entries) = std::fs::read_dir(&dir) else {
     76             continue;
     77         };
     78         for entry in entries.flatten() {
     79             let path = entry.path();
     80             match entry.file_type() {
     81                 Ok(kind) if kind.is_dir() => unvisited_dirs.push(path),
     82                 Ok(kind) if kind.is_file() && is_elf(&path) => elf_files.push(path),
     83                 _ => {}
     84             }
     85         }
     86     }
     87     elf_files.sort();
     88     elf_files
     89 }
     90 
     91 fn decode(path: &Path) -> Result<(load::Binary, Vec<arch::Ins>, HashMap<u64, String>)> {
     92     let bin = load::Binary::load(path)?;
     93     let cs = arch::build(bin.arch, bin.little)?;
     94     let insns = arch::sweep(&cs, &bin);
     95     let stubs = arch::resolve_stubs(&bin, &insns);
     96     Ok((bin, insns, stubs))
     97 }
     98 
     99 enum Skipped {
    100     UnsupportedArch(&'static str),
    101     Undecodable(anyhow::Error),
    102 }
    103 
    104 impl std::fmt::Display for Skipped {
    105     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    106         match self {
    107             Skipped::UnsupportedArch(arch) => {
    108                 write!(
    109                     f,
    110                     "{arch}, balrog only analyzes {}",
    111                     arch::ANALYZABLE_ARCHES
    112                 )
    113             }
    114             Skipped::Undecodable(err) => write!(f, "{err}"),
    115         }
    116     }
    117 }
    118 
    119 fn scan(path: &Path) -> Result<Vec<rules::Finding>, Skipped> {
    120     let (bin, insns, stubs) = decode(path).map_err(Skipped::Undecodable)?;
    121     if !bin.arch.analyzable() {
    122         return Err(Skipped::UnsupportedArch(bin.arch.name()));
    123     }
    124     if insns.is_empty() {
    125         return Ok(Vec::new());
    126     }
    127     let callgraph = graph::Graph::build(&bin, &insns, stubs);
    128     Ok(rules::analyze(&bin, &callgraph))
    129 }
    130 
    131 fn passes_filters(
    132     finding: &rules::Finding,
    133     min_score: i32,
    134     only_class: Option<rules::Class>,
    135 ) -> bool {
    136     finding.score >= min_score && only_class.is_none_or(|class| finding.class == class)
    137 }
    138 
    139 fn scan_and_filter(
    140     binaries: &[PathBuf],
    141     min_score: i32,
    142     only_class: Option<rules::Class>,
    143 ) -> Vec<rules::Finding> {
    144     let scan_results: Vec<Result<Vec<rules::Finding>, Skipped>> =
    145         binaries.par_iter().map(|binary| scan(binary)).collect();
    146     let mut findings = Vec::new();
    147     for (path, result) in binaries.iter().zip(scan_results) {
    148         match result {
    149             Ok(found) => findings.extend(
    150                 found
    151                     .into_iter()
    152                     .filter(|finding| passes_filters(finding, min_score, only_class)),
    153             ),
    154             Err(skipped) => eprintln!("warning: {}: {skipped}", path.display()),
    155         }
    156     }
    157     findings
    158 }
    159 
    160 fn dedup_templates(findings: &mut Vec<rules::Finding>) {
    161     let mut seen = HashSet::new();
    162     findings.retain(|finding| {
    163         let template = finding
    164             .evidence
    165             .iter()
    166             .find(|line| line.starts_with("format ") || line.starts_with("command template"))
    167             .cloned();
    168         seen.insert((
    169             finding.bin.clone(),
    170             finding.risky_call.clone(),
    171             finding.class,
    172             template,
    173         ))
    174     });
    175 }
    176 
    177 fn cap_per_binary(findings: &mut Vec<rules::Finding>, max_per_binary: usize) {
    178     let mut kept_so_far = HashMap::new();
    179     findings.retain(|finding| {
    180         let kept = kept_so_far
    181             .entry((finding.bin.clone(), finding.class))
    182             .or_insert(0usize);
    183         *kept += 1;
    184         *kept <= max_per_binary
    185     });
    186 }
    187 
    188 fn by_score(a: &rules::Finding, b: &rules::Finding) -> Ordering {
    189     b.score
    190         .cmp(&a.score)
    191         .then(a.bin.cmp(&b.bin))
    192         .then(a.site.cmp(&b.site))
    193         .then(a.risky_call.cmp(&b.risky_call))
    194 }
    195 
    196 fn by_class_then_score(a: &rules::Finding, b: &rules::Finding) -> Ordering {
    197     a.class.cmp(&b.class).then(by_score(a, b))
    198 }
    199 
    200 fn rank(mut findings: Vec<rules::Finding>, max_per_binary: usize) -> Vec<rules::Finding> {
    201     findings.sort_by(by_score);
    202     dedup_templates(&mut findings);
    203     cap_per_binary(&mut findings, max_per_binary);
    204     findings.sort_by(by_class_then_score);
    205     findings
    206 }
    207 
    208 fn print_finding(out: &mut impl Write, finding: &rules::Finding) -> io::Result<()> {
    209     writeln!(
    210         out,
    211         "\n[{:>3}] {}  {}+{:#x}  {}",
    212         finding.score, finding.risky_call, finding.bin, finding.site, finding.func
    213     )?;
    214     for line in &finding.evidence {
    215         writeln!(out, "       - {line}")?;
    216     }
    217     Ok(())
    218 }
    219 
    220 fn report(
    221     out: &mut impl Write,
    222     findings: &[rules::Finding],
    223     elf_count: usize,
    224     per_class_limit: usize,
    225 ) -> io::Result<()> {
    226     writeln!(out, "{elf_count} ELF, {} leads", findings.len())?;
    227     for class in rules::ORDER {
    228         let in_class: Vec<&rules::Finding> = findings
    229             .iter()
    230             .filter(|finding| finding.class == *class)
    231             .collect();
    232         if in_class.is_empty() {
    233             continue;
    234         }
    235         writeln!(
    236             out,
    237             "\n== {} ({} of {})",
    238             class.tag().to_uppercase(),
    239             in_class.len().min(per_class_limit),
    240             in_class.len()
    241         )?;
    242         for finding in in_class.iter().take(per_class_limit) {
    243             print_finding(out, finding)?;
    244         }
    245     }
    246     Ok(())
    247 }
    248 
    249 fn ignore_broken_pipe(r: io::Result<()>) -> Result<()> {
    250     match r {
    251         Err(e) if e.kind() == io::ErrorKind::BrokenPipe => Ok(()),
    252         other => Ok(other?),
    253     }
    254 }
    255 
    256 fn run_scan(args: ScanArgs) -> Result<()> {
    257     let binaries: Vec<PathBuf> = args
    258         .paths
    259         .iter()
    260         .flat_map(|root| collect_elf_files(root))
    261         .collect();
    262     let findings = scan_and_filter(&binaries, args.min_score, args.only_class);
    263     let ranked = rank(findings, args.max_per_binary);
    264 
    265     let mut out = BufWriter::new(io::stdout().lock());
    266     let written = if args.json {
    267         writeln!(out, "{}", serde_json::to_string_pretty(&ranked)?)
    268     } else {
    269         report(&mut out, &ranked, binaries.len(), args.per_class_limit)
    270     };
    271     ignore_broken_pipe(written.and_then(|_| out.flush()))
    272 }
    273 
    274 fn run_imports(path: &Path) -> Result<()> {
    275     let (bin, insns, stubs) = decode(path)?;
    276     if !bin.arch.analyzable() {
    277         anyhow::bail!(
    278             "{} is {}, balrog only analyzes {}",
    279             path.display(),
    280             bin.arch.name(),
    281             arch::ANALYZABLE_ARCHES
    282         );
    283     }
    284     println!(
    285         "{} {} {} insns={} jmprel={} undef={} stubs={}",
    286         bin.name(),
    287         bin.arch.name(),
    288         if bin.little { "LE" } else { "BE" },
    289         insns.len(),
    290         bin.got2sym.len(),
    291         bin.imports.len(),
    292         stubs.len()
    293     );
    294     let mut sorted_stubs: Vec<_> = stubs.iter().collect();
    295     sorted_stubs.sort();
    296     for (addr, name) in sorted_stubs.iter().take(STUB_PREVIEW) {
    297         println!("  {addr:#x} {name}");
    298     }
    299     Ok(())
    300 }
    301 
    302 fn main() -> Result<()> {
    303     match Cli::parse().cmd {
    304         Cmd::Scan(args) => run_scan(args),
    305         Cmd::Imports { path } => run_imports(&path),
    306     }
    307 }