Skip to content

Commit

Permalink
Adding basic vmlinux identification & extraction
Browse files Browse the repository at this point in the history
  • Loading branch information
rockrid3r committed Sep 14, 2023
1 parent e9aab76 commit 743d34e
Show file tree
Hide file tree
Showing 4 changed files with 200 additions and 18 deletions.
38 changes: 38 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ mod set_exec;
mod solvepy;
mod unstrip_libc;
mod warn;
mod proc_vmlinux;

pub use crate::pwninit::run;
pub use crate::pwninit::Result;
Expand Down Expand Up @@ -45,6 +46,26 @@ fn path_contains(path: &Path, pattern: &[u8]) -> bool {
.unwrap_or(false)
}

fn path_begins(path: &Path, prefix: &str) -> bool {
if let Some(filename) = path.file_name() {
if let Some(filename_str) = filename.to_str() {
filename_str.starts_with(prefix)
} else {
false
}
} else {
false
}
}

pub fn is_vmlinux(path: &Path) -> elf::detect::Result<bool> {
Ok(is_elf(path)? && path_begins(path, "vmlinux"))
}

pub fn is_bzimage(path: &Path) -> elf::detect::Result<bool> {
Ok(path_begins(path, "bzImage"))
}

/// Detect if `path` is the provided libc
pub fn is_libc(path: &Path) -> elf::detect::Result<bool> {
Ok(is_elf(path)? && path_contains(path, b"libc"))
Expand Down Expand Up @@ -104,6 +125,23 @@ pub fn set_bin_exec(opts: &Opts) -> io::Result<()> {
Ok(())
}

pub fn set_vmlinux_exec(opts: &Opts) -> io::Result<()> {
match &opts.vmlinux {
Some(vmlinux) => {
if !vmlinux.is_executable() {
println!(
"{}",
format!("setting {} executable", vmlinux.to_string_lossy().bold()).green()
);
set_exec(vmlinux)?;
}
},
None => "vmlinux not found".warn("failed setting vmlinux to be executable")
}

Ok(())
}

/// Set the detected linker executable
pub fn set_ld_exec(opts: &Opts) -> io::Result<()> {
match &opts.ld {
Expand Down
54 changes: 46 additions & 8 deletions src/opts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ use crate::is_bin;
use crate::is_ld;
use crate::is_libc;

use crate::is_vmlinux;
use crate::is_bzimage;

use std::path::Path;
use std::path::PathBuf;

Expand Down Expand Up @@ -78,11 +81,34 @@ pub struct Opts {
/// Disable generating template solve script
#[structopt(long)]
pub no_template: bool,

/// Kernel-land initialization
#[structopt(long)]
pub ker: bool,

/// Name of vmlinux binary
#[structopt(long)]
#[setters(generate)]
pub vmlinux: Option<PathBuf>,

/// Name of bzImage file
#[structopt(long)]
#[setters(generate)]
pub bzimage: Option<PathBuf>,

/// Disable extraction of vmlinux from bzImage
#[structopt(long)]
pub no_extract_vmlinux: bool,

/// Disable patching vmlinux with vmlinux-to-elf
#[structopt(long)]
pub no_patch_vmlinux: bool,
}

impl Opts {
/// Print the locations of known files (binary, libc, linker)
pub fn print(&self) {
println!("Found following {}-land challenge binaries:", if self.ker { "kernel" } else {"user"});
let f = |opt_path: &Option<PathBuf>, header: &str, color| {
if let Some(path) = opt_path {
println!(
Expand All @@ -93,9 +119,14 @@ impl Opts {
}
};

f(&self.bin, "bin", Color::BrightBlue);
f(&self.libc, "libc", Color::Yellow);
f(&self.ld, "ld", Color::Green);
if self.ker {
f(&self.bzimage, "bzImage", Color::BrightBlue);
f(&self.vmlinux, "vmlinux", Color::Green);
} else {
f(&self.bin, "bin", Color::BrightBlue);
f(&self.libc, "libc", Color::Yellow);
f(&self.ld, "ld", Color::Green);
}
}

/// For the unspecified files, try to guess their path
Expand All @@ -119,10 +150,17 @@ impl Opts {
Ok(if pred(&path)? { Some(path) } else { None })
};

Ok(self
.clone()
.with_bin(self.bin.or(f(is_bin)?))
.with_libc(self.libc.or(f(is_libc)?))
.with_ld(self.ld.or(f(is_ld)?)))
if self.ker {
Ok(self
.clone()
.with_bzimage(self.bzimage.or(f(is_bzimage)?))
.with_vmlinux(self.vmlinux.or(f(is_vmlinux)?)))
} else {
Ok(self
.clone()
.with_bin(self.bin.or(f(is_bin)?))
.with_libc(self.libc.or(f(is_libc)?))
.with_ld(self.ld.or(f(is_ld)?)))
}
}
}
80 changes: 80 additions & 0 deletions src/proc_vmlinux.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
use crate::opts::Opts;

use std::path::Path;
use std::process::Command;

use colored::Colorize;
use snafu::ResultExt;
use snafu::Snafu;
use std::fs::File;

use std::io;

#[derive(Debug, Snafu)]
#[allow(clippy::enum_variant_names)]
pub enum Error {
#[snafu(display("extract-vmlinux failed with nonzero exit status"))]
ExtractVmlinux,

#[snafu(display("extract-vmlinux failed to start; install it as shown in README.md"))]
ExtractVmlinuxExec { source: io::Error },

#[snafu(display("vmlinux-to-elf failed to start; install it as shown in README.md"))]
VmlinuxToElfExec { source: io::Error },

#[snafu(display("vmlinux-to-elf failed with nonzero exit status"))]
VmlinuxToElf,
}

pub type Result<T> = std::result::Result<T, Error>;

fn run_extract_vmlinux(bzimage: &Path) -> Result<()> {
println!(
"{}",
format!("running extract-vmlinux on {}", bzimage.to_string_lossy().bold()).green()
);

let mut cmd = Command::new("extract-vmlinux");
let new_vmlinux_file = File::create("./vmlinux").unwrap();
cmd.arg(bzimage).stdout(std::process::Stdio::from(new_vmlinux_file));
let status = cmd.status().context(ExtractVmlinuxExecSnafu)?;
if status.success() {
Ok(())
} else {
Err(Error::ExtractVmlinux)
}
}

fn run_vmlinux_to_elf(vmlinux: &Path) -> Result<()> {
println!(
"{}",
format!("running vmlinux-to-elf on {}", vmlinux.to_string_lossy().bold()).green()
);

let mut cmd = Command::new("vmlinux-to-elf");
cmd
.arg(vmlinux)
.arg(format!("{}_patched", vmlinux.to_string_lossy()));
let status = cmd.status().context(VmlinuxToElfExecSnafu)?;
if status.success() {
Ok(())
} else {
Err(Error::VmlinuxToElf)
}
}

pub fn extract_vmlinux_from_bzimage(opts: &Opts) -> Result<()> {
if let Some(bzimage) = &opts.bzimage {
run_extract_vmlinux(bzimage)?;
}

Ok(())
}

pub fn patch_vmlinux(opts: &Opts) -> Result<()> {
if let Some(vmlinux) = &opts.vmlinux {
run_vmlinux_to_elf(vmlinux)?;
}

Ok(())
}
46 changes: 36 additions & 10 deletions src/pwninit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@ use crate::opts;
use crate::patch_bin;
use crate::set_bin_exec;
use crate::set_ld_exec;
use crate::set_vmlinux_exec;
use crate::solvepy;
use crate::Opts;
use crate::proc_vmlinux;

use ex::io;
use snafu::ResultExt;
Expand All @@ -28,6 +30,15 @@ pub enum Error {

#[snafu(display("failed making template solve script: {}", source))]
Solvepy { source: solvepy::Error },

#[snafu(display("failed extracting vmlinux: {}", source))]
ExtractVmlinux { source: proc_vmlinux::Error },

#[snafu(display("failed setting vmlinux executable: {}", source))]
SetVmlinuxExec { source: io::Error },

#[snafu(display("failed patching vmlinux: {}", source))]
PatchVmlinux { source: proc_vmlinux::Error },
}

pub type Result = std::result::Result<(), Error>;
Expand All @@ -41,20 +52,35 @@ pub fn run(opts: Opts) -> Result {
opts.print();
println!();

set_bin_exec(&opts).context(SetBinExecSnafu)?;
maybe_visit_libc(&opts);
if opts.ker {
if !opts.no_extract_vmlinux {
proc_vmlinux::extract_vmlinux_from_bzimage(&opts).context(ExtractVmlinuxSnafu)?;
}

// Redo detection in case the ld was downloaded
let opts = opts.find_if_unspec().context(FindSnafu)?;
// Redo detection in case the vmlinux was extracted
let opts = opts.find_if_unspec().context(FindSnafu)?;

set_ld_exec(&opts).context(SetLdExecSnafu)?;
if !opts.no_patch_vmlinux {
proc_vmlinux::patch_vmlinux(&opts).context(PatchVmlinuxSnafu)?;
}

if !opts.no_patch_bin {
patch_bin::patch_bin(&opts).context(PatchBinSnafu)?;
}
set_vmlinux_exec(&opts).context(SetVmlinuxExecSnafu)?;
} else {
set_bin_exec(&opts).context(SetBinExecSnafu)?;
maybe_visit_libc(&opts);

// Redo detection in case the ld was downloaded
let opts = opts.find_if_unspec().context(FindSnafu)?;

set_ld_exec(&opts).context(SetLdExecSnafu)?;

if !opts.no_patch_bin {
patch_bin::patch_bin(&opts).context(PatchBinSnafu)?;
}

if !opts.no_template {
solvepy::write_stub(&opts).context(SolvepySnafu)?;
if !opts.no_template {
solvepy::write_stub(&opts).context(SolvepySnafu)?;
}
}

Ok(())
Expand Down

0 comments on commit 743d34e

Please sign in to comment.