From 2889470031a24500d5b3e339980bf450d7a761a6 Mon Sep 17 00:00:00 2001
From: alemi <me@alemi.dev>
Date: Mon, 3 Apr 2023 05:17:59 +0200
Subject: [PATCH] feat: added some commands to lua shell

---
 Cargo.toml |   1 +
 src/lib.rs | 110 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 105 insertions(+), 6 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2a9e667..d19fb1e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,3 +19,4 @@ tracing-subscriber = "0.3"
 mlua = { version = "0.8", features = ["luajit52", "vendored", "async", "send", "serialize"] }
 serde = "1.0.159"
 serde_json = "1.0.95"
+hexdump = "0.1.1"
diff --git a/src/lib.rs b/src/lib.rs
index f12d3a1..03504cd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,9 @@
-use mlua::{Lua, MultiValue};
+use mlua::{Lua, MultiValue, Variadic, Value, Table};
 use tokio::{sync::{mpsc, broadcast}, net::{TcpStream, TcpListener}, io::{AsyncWriteExt, AsyncReadExt}};
-use tracing::{error, debug, info};
+use tracing::{error, debug};
+
+use pox::proc_maps::get_process_maps;
+use pox::tricks::fmt_path;
 
 #[ctor::ctor]
 fn contructor() {
@@ -21,7 +24,6 @@ fn contructor() {
 fn destructor() {}
 
 async fn main() {
-
 	let mut handle = ControlChannel::run("127.0.0.1:13337".into());
 
 	loop {
@@ -83,7 +85,8 @@ impl ControlChannel {
 	}
 
 	async fn process(&mut self, mut stream: TcpStream) {
-		let lua = Lua::new();
+		let mut lua = Lua::new();
+		prepare_lua_runtime(&mut lua, self.source.clone());
 		self.source.send(
 			format!("LuaJit 5.2 via rlua inside process #{}\n@> ", std::process::id())
 		).unwrap();
@@ -108,8 +111,7 @@ impl ControlChannel {
 								match lua.load(&cmd).eval::<MultiValue>() {
 									Ok(values) => {
 										for val in values {
-											let x = serde_json::to_string(&val).unwrap();
-											self.source.send(format!("@({}) : {}",val.type_name(), x)).unwrap();
+											self.source.send(format!("=({}) {}", val.type_name(), pretty_lua(val))).unwrap();
 										}
 										self.source.send("\n@> ".into()).unwrap();
 										cmd = String::new();
@@ -149,3 +151,99 @@ impl ControlChannel {
 		}
 	}
 }
+
+fn prepare_lua_runtime(lua: &Lua, console: broadcast::Sender<String>) {
+	let c = console.clone();
+	let log = lua.create_function(move |_lua, values: Variadic<Value>| {
+		let mut out = String::new();
+		for value in values {
+			out.push_str(&pretty_lua(value));
+			out.push(' ');
+		}
+		out.push('\n');
+		let size = out.len();
+		c.send(out).unwrap();
+		Ok(size)
+	}).unwrap();
+	lua.globals().set("log", log).unwrap();
+
+	let procmaps = lua.create_function(move |_lua, ()| {
+		let mut out = String::new();
+		for map in get_process_maps(std::process::id() as i32).unwrap() {
+			out.push_str(
+				format!(
+					"[{}] 0x{:08X}..0x{:08X} +{:08x} \t {} {}\n",
+					map.flags, map.start(), map.start() + map.size(), map.offset, fmt_path(map.filename()),
+					if map.inode != 0 { format!("({})", map.inode) } else { "".into() },
+				).as_str()
+			);
+		}
+		Ok(out)
+	}).unwrap();
+	lua.globals().set("procmaps", procmaps).unwrap();
+
+	let hexdump = lua.create_function(move |_lua, (addr, size): (usize, usize)| {
+		if size == 0 {
+			return Ok("".into());
+		}
+		let ptr = addr as *mut u8;
+		let slice = unsafe { std::slice::from_raw_parts(ptr, size) };
+		let mut out = String::new();
+		for line in hexdump::hexdump_iter(slice) {
+			out.push_str(&line);
+			out.push('\n');
+		}
+		Ok(out)
+	}).unwrap();
+	lua.globals().set("hexdump", hexdump).unwrap();
+
+	let write = lua.create_function(move |_lua, (addr, data): (usize, Vec<u8>)| {
+		for (i, byte) in data.iter().enumerate() {
+			let off = (addr + i) as *mut u8;
+			unsafe { *off = *byte } ;
+		}
+		Ok(data.len())
+	}).unwrap();
+	lua.globals().set("write", write).unwrap();
+
+	let exit = lua.create_function(move |_lua, code: i32| {
+		#[allow(unreachable_code)]
+		Ok(std::process::exit(code))
+	}).unwrap();
+	lua.globals().set("exit", exit).unwrap();
+
+	let help = lua.create_function(move |_lua, ()| {
+		console.send(" > log(...)              print to (this) remote shell\n".into()).unwrap();
+		console.send(" > exit(code)            immediately terminate process\n".into()).unwrap();
+		console.send(" > procmaps()            returns process memory maps as string\n".into()).unwrap();
+		console.send(" > write(addr, bytes)    write raw bytes at given address\n".into()).unwrap();
+		console.send(" > hexdump(addr, size)   dump bytes at addr in hexdump format\n".into()).unwrap();
+		console.send(" > help()                print these messages".into()).unwrap();
+		Ok(())
+	}).unwrap();
+	lua.globals().set("help", help).unwrap();
+}
+
+fn pretty_lua(val: Value) -> String {
+	// TODO there must be some builtin to do this, right???
+	match val {
+		Value::Nil => "nil".into(),
+		Value::Boolean(b) => if b { "true".into() } else { "false".into() },
+		Value::LightUserData(x) => format!("LightUserData({:?})", x),
+		Value::Integer(n) => format!("{}", n),
+		Value::Number(n) => format!("{:.3}", n),
+		Value::String(s) => s.to_str().expect("string is not str").into(),
+		Value::Table(t) => try_serialize_table(&t),
+		Value::Function(f) => format!("Function({:?}", f),
+		Value::Thread(t) => format!("Thread({:?})", t),
+		Value::UserData(x) => format!("UserData({:?})", x),
+		Value::Error(e) => format!("Error({:?}) : {}", e, e.to_string()),
+	}
+}
+
+fn try_serialize_table(t: &Table) -> String {
+	match serde_json::to_string(t) {
+		Ok(txt) => txt,
+		Err(_e)  => format!("{:?}", t),
+	}
+}