Implement puppeting

client/src/proc.rs

@@ -0,0 +1,273 @@
+// Full credit to https://github.com/ohchase/ptrace_syscalls/tree/master
+
+use std::vec;
+
+use nix::{
+    libc::user_regs_struct,
+    sys::{
+        ptrace,
+        signal::Signal::SIGTRAP,
+        wait::{WaitStatus, waitpid},
+    },
+    unistd::Pid,
+};
+use syscalls::Sysno;
+
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum HostError {
+    #[error("Process not found `{0}`")]
+    ProcessNotFound(String),
+    #[error("Nix Error `{0}`")]
+    NixError(#[from] nix::errno::Errno),
+    #[error("Unexpected Wait Status `{0:#?}`")]
+    UnexpectedWaitStatus(WaitStatus),
+    #[error("StdIo Error `{0}`")]
+    Io(#[from] std::io::Error),
+    #[error("Mmap Error `{0:#?}`")]
+    MmapBadAddress(u64),
+    #[error("Munmap Error `{0:#?}`")]
+    MunmapFailed(u64),
+}
+
+pub type HostResult<T> = Result<T, HostError>;
+
+pub struct UserProcessMemory<'a> {
+    address: u64,
+    owner: &'a UserProcess,
+    len: u64,
+}
+
+impl<'a> UserProcessMemory<'a> {
+    /// Getter for UserProcessMemory's start address
+    pub fn address(&self) -> u64 {
+        self.address
+    }
+
+    /// Getter for UserProcessMemory's length
+    #[allow(clippy::len_without_is_empty)]
+    pub fn len(&self) -> u64 {
+        self.len
+    }
+}
+
+impl<'a> Drop for UserProcessMemory<'a> {
+    /// Munmap's the memory on drop
+    fn drop(&mut self) {
+        match self.owner.deallocate_memory(self.address, self.len) {
+            Ok(()) => log::trace!(
+                "Successfully deallocated memory {:#X} with length {}",
+                self.address,
+                self.len
+            ),
+            Err(err) => log::error!(
+                "Failed to deallocate memory {:#X} with length {}, Error: {:#?}",
+                self.address,
+                self.len,
+                err
+            ),
+        }
+    }
+}
+
+pub struct UserProcess {
+    pid: Pid,
+}
+
+impl UserProcess {
+    /// Initializes the UserProcess
+    /// Given you have a UserProcess instance it means the attach must of succeeded.
+    /// Attach should be the only way to acquire a UserProcess.
+    pub fn attach(pid: Pid) -> HostResult<Self> {
+        ptrace::attach(pid)?;
+
+        log::info!("New UserProcess successfully attached to pid: {}", pid);
+        Ok(Self { pid })
+    }
+
+    /// Getter for UserProcess actively connected pid
+    pub fn pid(&self) -> Pid {
+        self.pid
+    }
+
+    /// Uses mmap syscall with ptrace to allocate user process memory
+    pub fn allocate_memory(
+        &self,
+        address: u64,
+        len: u64,
+        prot: u64,
+        flags: u64,
+        fd: u64,
+        offset: u64,
+    ) -> HostResult<UserProcessMemory> {
+        let mmap_result =
+            self.sys_call(Sysno::mmap as u64, address, len, prot, flags, fd, offset)?;
+        let mmap_result = mmap_result.rax;
+
+        // invalid address
+        // TODO not quite right...
+        if mmap_result == 0 {
+            return Err(HostError::MmapBadAddress(address));
+        }
+
+        Ok(UserProcessMemory {
+            address: mmap_result,
+            owner: self,
+            len,
+        })
+    }
+
+    /// Uses munmap syscall with ptrace to deallocate user process memory
+    fn deallocate_memory(&self, address: u64, len: u64) -> HostResult<()> {
+        let munmap_result = self.sys_call(Sysno::munmap as u64, address, len, 0, 0, 0, 0)?;
+        let munmap_result = munmap_result.rax;
+
+        // not a zero return value means a failure
+        if munmap_result != 0 {
+            return Err(HostError::MunmapFailed(address));
+        }
+
+        Ok(())
+    }
+
+    /// Write to the user process memory
+    pub fn write_user_memory(
+        &self,
+        user_memory: &mut UserProcessMemory,
+        offset: u64,
+        bytes: &[u8],
+    ) -> HostResult<usize> {
+        self.write_memory(user_memory.address + offset, bytes)
+    }
+
+    /// Read from user process memory
+    pub fn read_user_memory(
+        &self,
+        user_memory: &UserProcessMemory,
+        len: usize,
+    ) -> HostResult<Vec<u8>> {
+        self.read_memory(user_memory.address, len)
+    }
+
+    /// String to the proc's memory file
+    /// Reference: https://crates.io/crates/pete
+    fn proc_mem_path(&self) -> String {
+        format!("/proc/{}/mem", self.pid.as_raw() as u32)
+    }
+
+    /// Common wrapper around writing memory
+    /// Returns the amount of bytes written
+    /// Reference: https://crates.io/crates/pete
+    fn write_memory(&self, addr: u64, data: &[u8]) -> HostResult<usize> {
+        use std::os::unix::fs::FileExt;
+        let mem = std::fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(self.proc_mem_path())?;
+        let len = mem.write_at(data, addr)?;
+        Ok(len)
+    }
+
+    /// Common wrapper around reading memory
+    /// Returns a vector to the memory bytes
+    /// Reference: https://crates.io/crates/pete
+    fn read_memory(&self, addr: u64, len: usize) -> HostResult<Vec<u8>> {
+        use std::os::unix::fs::FileExt;
+
+        let mut data = vec![0u8; len];
+        let mem = std::fs::File::open(self.proc_mem_path())?;
+        let len_read = mem.read_at(&mut data, addr)?;
+
+        data.truncate(len_read);
+        Ok(data)
+    }
+
+    /// Invokes a syscall in the userprocess
+    /// Accepts up to six arguments
+    #[allow(clippy::too_many_arguments)]
+    pub fn sys_call(
+        &self,
+        sys_call: u64,
+        rdi: u64,
+        rsi: u64,
+        rdx: u64,
+        r10: u64,
+        r8: u64,
+        r9: u64,
+    ) -> HostResult<user_regs_struct> {
+        log::trace!("UserProcess {} Syscall: {:#?}", self.pid, sys_call);
+        let syscall_instruction = [0x0Fu8, 0x05u8];
+
+        // Cache original registers, original instruction pointer (rip), and the original instructions
+        let original_registers = ptrace::getregs(self.pid)?;
+        let original_ip = original_registers.rip;
+        let original_instructions = self.read_memory(original_ip, syscall_instruction.len())?;
+
+        // Write over our shell code 0x0F05 for sys call
+        self.write_memory(original_ip, &syscall_instruction)?;
+
+        // Create a copy of the original registers
+        // Set the sys_call index and args[0..5] (six arguments)
+        let mut new_registers = original_registers;
+        new_registers.rax = sys_call;
+        new_registers.rdi = rdi;
+        new_registers.rsi = rsi;
+        new_registers.rdx = rdx;
+        new_registers.r10 = r10;
+        new_registers.r8 = r8;
+        new_registers.r9 = r9;
+
+        // Apply the new registers, and new instructions then single step waiting for SIG_TRAP
+        ptrace::setregs(self.pid, new_registers)?;
+
+        // Single step the process and wait for a SIGTRAP signal
+        self.single_step()?;
+
+        // Cache the resultant registers
+        let result = ptrace::getregs(self.pid)?;
+
+        // Restore original instructions, and original registers to continue normal program control flow
+        self.write_memory(original_ip, &original_instructions)?;
+        ptrace::setregs(self.pid, original_registers)?;
+
+        Ok(result)
+    }
+
+    /// Single steps the process
+    /// Fails if the next signal is not a SIGTRAP
+    fn single_step(&self) -> HostResult<()> {
+        ptrace::step(self.pid, None)?;
+        self.wait_trap()
+    }
+
+    /// Waits for the next signal
+    /// Returns an error if that signal is not a SIGTRAP
+    fn wait_trap(&self) -> HostResult<()> {
+        match self.wait()? {
+            WaitStatus::Stopped(_, SIGTRAP) => Ok(()),
+            status => Err(HostError::UnexpectedWaitStatus(status)),
+        }
+    }
+
+    /// Waits on the current process
+    fn wait(&self) -> HostResult<WaitStatus> {
+        waitpid(self.pid, None).map_err(HostError::NixError)
+    }
+}
+
+impl Drop for UserProcess {
+    /// Drop implementation for the UserProcess
+    /// Attempts to detach from the process, does not panic on error instead only logs
+    fn drop(&mut self) {
+        if let Err(err) = ptrace::detach(self.pid, None) {
+            log::error!(
+                "UserProcess failed to detach from: {}, with Err: {:#?}",
+                self.pid,
+                err
+            );
+        } else {
+            log::trace!("UserProcess successfully detached from: {}", self.pid);
+        }
+    }
+}