Adding basic vmlinux identification & extraction

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<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"))
@@ -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 {

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<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!(
@@ -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)?)))
+        }
     }
 }

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<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(())
+}

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(())