chore: refactored a lot, made it fancy

src/needle/main.rs

@@ -1,9 +1,10 @@
 mod syscalls;
+mod operations;
 
-use std::ffi::c_void;
+use nix::{Result, {sys::{ptrace, wait::waitpid}, unistd::Pid}};
-use nix::{Result, {sys::{ptrace, wait::waitpid}, unistd::Pid}, libc::{PROT_READ, PROT_WRITE, MAP_PRIVATE, MAP_ANON}};
 use clap::Parser;
-use syscalls::{prepare_mmap, prepare_write};
+use operations::step_to_syscall;
+use syscalls::{RemoteString, RemoteWrite, RemoteSyscall};
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -16,97 +17,12 @@ struct NeedleArgs {
 	word: u32,
 }
 
-pub fn send_str(pid: Pid, syscall_addr: u64, data: &str) -> Result<usize> {
+pub fn nasty_stuff(pid: Pid, _word_size: usize) -> Result<()> {
-	let mut regs = ptrace::getregs(pid)?;
+	let original_registers = ptrace::getregs(pid)?;
-	regs.rip = syscall_addr;
+	let syscall_addr = step_to_syscall(pid)?;
-	prepare_mmap(&mut regs, 0, data.len(), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, 0xFFFFFFFFFFFFFFFF, 0);
+	let msg = RemoteString::new(pid, syscall_addr, "injected!\n\0".into())?;
-	ptrace::setregs(pid, regs)?;
+	RemoteWrite::args(1, msg).syscall(pid, syscall_addr)?;
-	ptrace::step(pid, None)?;
+	ptrace::setregs(pid, original_registers)?;
-	waitpid(pid, None)?;
-	let zone_addr = ptrace::getregs(pid)?.rax as usize;
-	write_buffer(pid, zone_addr, data.as_bytes(), 32)?; // TODO word size!
-	Ok(zone_addr)
-}
-
-pub fn read_buffer(pid: Pid, addr: usize, size: usize, word: u32) -> Result<Vec<u8>> {
-	let mut out = vec![];
-
-	for i in (0..size).step_by((word/8) as usize) {
-		let data = ptrace::read(pid, (addr + i) as *mut c_void)?;
-		for j in 0..(word/8) as usize {
-			out.push(((data >> (j * 8)) & 0xFF) as u8);
-		}
-	}
-
-	Ok(out)
-}
-
-pub fn write_buffer(pid: Pid, addr: usize, payload: &[u8], word:u32) -> Result<()> {
-	let step = word / 8;
-	let mut at = addr;
-
-	for chunk in payload.chunks(step as usize) {
-		let mut buf : u64 = 0;
-		for (i, c) in chunk.iter().enumerate() {
-			buf |= (*c as u64) << (i * 8);
-		}
-		unsafe { ptrace::write(pid, at as *mut c_void, buf as *mut c_void)?; }
-		at += step as usize;
-	}
-
-	Ok(())
-}
-
-pub fn pwn(pid: Pid, _word_size: usize) -> Result<()> {
-
-	let mut prev_regs;
-	let mut insn_addr;
-	let mut curr_instr;
-
-	// seek to syscall
-	loop {
-		prev_regs = ptrace::getregs(pid)?;
-		insn_addr = prev_regs.rip;
-		curr_instr = ptrace::read(pid, insn_addr as *mut c_void)?;
-		// println!("@ 0x{:X} [{:x}]", insn_addr, curr_instr);
-
-		if curr_instr & 0xFFFF == 0x050F {
-			// println!("found syscall!");
-			break;
-		}
-
-		ptrace::step(pid, None)?;
-		waitpid(pid, None)?;
-	}
-
-	// // Put syscall opcodes
-	// let mut syscall_insn = vec![0x00u8; word_size/8];
-	// syscall_insn[0] = 0x05; // it's two!
-	// syscall_insn[1] = 0x0F;
-
-	// unsafe {
-	// 	ptrace::write(pid, insn_addr, syscall_insn.as_slice().as_ptr() as *mut c_void)?;
-	// }
-
-	let msg = send_str(pid, insn_addr, "injected!\n\0")?;
-
-	let mut call_regs = prev_regs.clone();
-
-	call_regs.rip = insn_addr;
-
-	prepare_write(&mut call_regs, 1, msg, 10);
-	ptrace::setregs(pid, call_regs)?;
-	ptrace::step(pid, None)?;
-	waitpid(pid, None)?;
-	// println!("Written payload to stdout on tracee");
-
-	// restore code and registers
-	// unsafe {
-	// 	ptrace::write(pid, insn_addr, prev_instr.as_ptr() as *mut c_void)?;
-	// }
-
-	ptrace::setregs(pid, prev_regs)?;
-
 	Ok(())
 }
 
@@ -125,7 +41,7 @@ fn main() {
 
 	println!("Attached to process #{}", args.pid);
 
-	if let Err(e) = pwn(pid, args.word as usize) {
+	if let Err(e) = nasty_stuff(pid, args.word as usize) {
 		eprintln!("Could not pwn : {}", e);
 	}
 

src/needle/operations.rs

@@ -0,0 +1,54 @@
+use std::ffi::c_void;
+
+use nix::{Result, unistd::Pid, sys::{ptrace, wait::waitpid}};
+
+#[allow(unused)]
+pub fn read_buffer(pid: Pid, addr: usize, size: usize, word: u32) -> Result<Vec<u8>> {
+	let mut out = vec![];
+
+	for i in (0..size).step_by((word/8) as usize) {
+		let data = ptrace::read(pid, (addr + i) as *mut c_void)?;
+		for j in 0..(word/8) as usize {
+			out.push(((data >> (j * 8)) & 0xFF) as u8);
+		}
+	}
+
+	Ok(out)
+}
+
+pub fn write_buffer(pid: Pid, addr: usize, payload: &[u8], word:u32) -> Result<()> {
+	let step = word / 8;
+	let mut at = addr;
+
+	for chunk in payload.chunks(step as usize) {
+		let mut buf : u64 = 0;
+		for (i, c) in chunk.iter().enumerate() {
+			buf |= (*c as u64) << (i * 8);
+		}
+		unsafe { ptrace::write(pid, at as *mut c_void, buf as *mut c_void)?; }
+		at += step as usize;
+	}
+
+	Ok(())
+}
+
+pub fn step_to_syscall(pid: Pid) -> Result<usize> {
+	let mut registers;
+	let mut addr;
+	let mut instructions;
+
+	// seek to syscall
+	loop {
+		registers = ptrace::getregs(pid)?;
+		addr = registers.rip as usize;
+		instructions = ptrace::read(pid, addr as *mut c_void)?;
+		// println!("@ 0x{:X} [{:x}]", insn_addr, curr_instr);
+
+		if instructions & 0xFFFF == 0x050F {
+			return Ok(addr);
+		}
+
+		ptrace::step(pid, None)?;
+		waitpid(pid, None)?;
+	}
+}

src/needle/syscalls.rs

@@ -1,31 +1,117 @@
-use nix::libc::user_regs_struct;
+use nix::{libc::{user_regs_struct, MAP_PRIVATE, MAP_ANON, PROT_READ, PROT_WRITE}, Result, sys::{ptrace, wait::waitpid}, unistd::Pid};
 
-fn prepare_for_syscall(regs: &mut user_regs_struct, rax: u64, rdi: u64, rsi: u64, rdx: u64, r10: u64, r8: u64, r9: u64) {
+use crate::operations::write_buffer;
-	regs.rax = rax;
+
-	regs.rdi = rdi;
+pub struct RemoteString {
-	regs.rsi = rsi;
+	pub ptr: usize,
-	regs.rdx = rdx;
+	pub txt: String,
-	regs.r10 = r10;
+}
-	regs.r8  = r8;
+
-	regs.r9  = r9;
+impl RemoteString {
+	pub fn new(pid: Pid, syscall: usize, txt: String) -> Result<Self> {
+		let ptr = RemoteMMap::args(
+			0, txt.len(), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, 0xFFFFFFFFFFFFFFFF, 0
+		).syscall(pid, syscall)? as usize;
+		write_buffer(pid, ptr, txt.as_bytes(), 32)?; // TODO don't hardcode word size
+		Ok(RemoteString { ptr, txt })
+	}
+}
+
+pub trait RemoteSyscall {
+	fn registers(&self, regs: &mut user_regs_struct);
+
+	fn syscall(&self, pid: Pid, addr: usize) -> Result<u64> {
+		let mut regs = ptrace::getregs(pid)?;
+		regs.rip = addr as u64;
+		self.registers(&mut regs);
+		ptrace::setregs(pid, regs)?;
+		ptrace::step(pid, None)?;
+		waitpid(pid, None)?;
+		regs = ptrace::getregs(pid)?;
+		Ok(regs.rax)
+	}
+
+	fn prepare_registers(regs: &mut user_regs_struct, rax: u64, rdi: u64, rsi: u64, rdx: u64, r10: u64, r8: u64, r9: u64) {
+		regs.rax = rax;
+		regs.rdi = rdi;
+		regs.rsi = rsi;
+		regs.rdx = rdx;
+		regs.r10 = r10;
+		regs.r8  = r8;
+		regs.r9  = r9;
+	}
+}
+
+pub struct RemoteMMap {
+	addr: u64,
+	len: usize,
+	prot: i32,
+	flags: i32,
+	fd: u64,
+	off: u64,
+}
+
+impl RemoteMMap {
+	pub fn args(addr: u64, len: usize, prot: i32, flags: i32, fd: u64, off: u64) -> Self {
+		RemoteMMap { addr, len, prot, flags, fd, off }
+	}
+}
+
+impl RemoteSyscall for RemoteMMap {
+	fn registers(&self, regs: &mut user_regs_struct) {
+		Self::prepare_registers(regs, 9, self.addr, self.len as u64, self.prot as u64, self.flags as u64, self.fd, self.off);
+	}
+}
+
+pub struct RemoteOpen {
+	filename: RemoteString,
+	flags: u64,
+	mode: u64,
 }
 
 #[allow(unused)]
-pub fn prepare_mmap(regs: &mut user_regs_struct, addr: u64, len: usize, prot: i32, flags: i32, fd: u64, off: u64) {
+impl RemoteOpen {
-	prepare_for_syscall(regs, 9, addr, len as u64, prot as u64, flags as u64, fd, off);
+	pub fn args(filename: RemoteString, flags: u64, mode: u64) -> Self {
+		RemoteOpen { filename, flags, mode }
+	}
+}
+
+impl RemoteSyscall for RemoteOpen {
+	fn registers(&self, regs: &mut user_regs_struct) {
+		Self::prepare_registers(regs, 2, self.filename.ptr as u64, self.flags, self.mode, 0, 0, 0);
+	}
+}
+
+pub struct RemoteWrite {
+	fd: u64,
+	buf: RemoteString, // TODO make remote slice or remote bytes
+}
+
+impl RemoteWrite {
+	pub fn args(fd: u64, buf: RemoteString) -> Self {
+		RemoteWrite { fd, buf }
+	}
+}
+
+impl RemoteSyscall for RemoteWrite {
+	fn registers(&self, regs: &mut user_regs_struct) {
+		Self::prepare_registers(regs, 1, self.fd, self.buf.ptr as u64, self.buf.txt.len() as u64, 0, 0, 0);
+	}
+}
+
+pub struct RemoteExit {
+	code: i64,
 }
 
 #[allow(unused)]
-pub fn prepare_open(regs: &mut user_regs_struct, filename: &str, flags: u64, mode: u64) {
+impl RemoteExit {
-	prepare_for_syscall(regs, 2, filename.as_ptr() as u64, flags, mode, 0, 0, 0);
+	pub fn args(code: i64) -> Self {
+		RemoteExit { code }
+	}
 }
 
-#[allow(unused)]
+impl RemoteSyscall for RemoteExit {
-pub fn prepare_write(regs: &mut user_regs_struct, fd: u64, buf: usize, count: usize) {
+	fn registers(&self, regs: &mut user_regs_struct) {
-	prepare_for_syscall(regs, 1, fd, buf as u64, count as u64, 0, 0, 0);
+	 Self::prepare_registers(regs, 60, self.code as u64, 0, 0, 0, 0, 0);
-}
+	}
-
-#[allow(unused)]
-pub fn prepare_exit(regs: &mut user_regs_struct, error_code: i64) {
-	prepare_for_syscall(regs, 60, error_code as u64, 0, 0, 0, 0, 0);
 }