diff --git a/src/lib.rs b/src/lib.rs index 0890803..24762c1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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; @@ -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 { + Ok(is_elf(path)? && path_begins(path, "vmlinux")) +} + +pub fn is_bzimage(path: &Path) -> elf::detect::Result { + Ok(path_begins(path, "bzImage")) +} + /// Detect if `path` is the provided libc pub fn is_libc(path: &Path) -> elf::detect::Result { Ok(is_elf(path)? && path_contains(path, b"libc")) @@ -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 { diff --git a/src/opts.rs b/src/opts.rs index a6be341..76c9f8d 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -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; @@ -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, + + /// Name of bzImage file + #[structopt(long)] + #[setters(generate)] + pub bzimage: Option, + + /// 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, header: &str, color| { if let Some(path) = opt_path { println!( @@ -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 @@ -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)?))) + } } } diff --git a/src/proc_vmlinux.rs b/src/proc_vmlinux.rs new file mode 100644 index 0000000..5ce07cc --- /dev/null +++ b/src/proc_vmlinux.rs @@ -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 = std::result::Result; + +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(()) +} diff --git a/src/pwninit.rs b/src/pwninit.rs index e518abf..0d89499 100644 --- a/src/pwninit.rs +++ b/src/pwninit.rs @@ -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; @@ -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>; @@ -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(())