// SPDX-FileCopyrightText: edef <edef@unfathomable.blue>
// SPDX-FileCopyrightText: V <v@unfathomable.blue>
// SPDX-License-Identifier: OSL-3.0

use {
	anyhow::{bail, Context, Result},
	bitflags::bitflags,
	nix::{
		libc,
		sys::{
			personality::{self, Persona},
			ptrace,
			signal::Signal,
			wait::{waitpid, WaitPidFlag, WaitStatus},
		},
		unistd::Pid,
	},
	std::{
		convert::TryInto,
		env,
		ffi::CString,
		fs::File,
		io::{self, BufRead, Seek, SeekFrom},
		os::unix::process::CommandExt,
		process::Command,
	},
};

// TODO(edef): consider implementing this in terms of TID?
// tgids are a strict subset of tids
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Tgid(pub libc::pid_t);

impl Tgid {
	fn as_pid(&self) -> Pid {
		Pid::from_raw(self.0)
	}
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Tid(pub libc::pid_t);

impl Tid {
	fn as_pid(&self) -> Pid {
		Pid::from_raw(self.0)
	}
}

#[derive(Debug)]
struct Process {
	tgid: Tgid,
	mem: File,
}

impl Process {
	fn spawn(cmd: &mut Command) -> Result<Process> {
		unsafe {
			cmd.pre_exec(|| {
				// disable ASLR
				let mut persona = personality::get()?;
				persona.insert(Persona::ADDR_NO_RANDOMIZE);
				personality::set(persona)?;

				ptrace::traceme()?;
				Ok(())
			});
		}

		let child = cmd.spawn()?;

		// the thread group leader's TID is equal to the TGID
		let tgid = Tgid(child.id() as _);

		match waitpid(tgid.as_pid(), None).context("Couldn't waitpid on fresh child")? {
			WaitStatus::Stopped(_, Signal::SIGTRAP) => {}
			status => bail!("unexpected child state: {status:?}"),
		}

		Ok(Process {
			tgid,
			mem: File::open(format!("/proc/{}/mem", tgid.0))
				.context("Couldn't open child memory")?,
		})
	}

	fn read_mem_cstr(&self, ptr: u64) -> Result<CString> {
		let mut mem = io::BufReader::new(&self.mem);
		mem.seek(SeekFrom::Start(ptr))?;
		let mut buf = vec![];
		mem.read_until(0, &mut buf)?;

		Ok(CString::from_vec_with_nul(buf).expect("logic error"))
	}
}

macro_rules! define_syscalls {
	(enum $SyscallEntry:ident {
		$(fn $syscall:ident ( $($arg:ident : $Arg:ty),* ) -> $Ret:ty = $nr:literal ;)*
	}) => {
		#[derive(Debug, Clone)]
		#[allow(non_camel_case_types)]
		// TODO(edef): re-enable dead_code lint when we start fully interpreting syscall args
		#[allow(dead_code)]
		enum $SyscallEntry {
			$($syscall {
				$($arg : $Arg),*
			}),*
		}

		impl $SyscallEntry {
			fn from_regs(regs: libc::user_regs_struct) -> Result<$SyscallEntry> {
				Ok(match (regs.orig_rax, [regs.rdi, regs.rsi, regs.rdx, regs.r10, regs.r8, regs.r9]) {
					$(
						($nr, [$($arg),*, ..]) => $SyscallEntry::$syscall {
							$($arg: match SyscallArg::try_from_reg($arg) {
								Some(x) => x,
								None => bail!("couldn't parse {}(2) {}: 0x{:08x}", stringify!($syscall), stringify!($arg), $arg)
							}),*
						},
					)*
					(n, _) => bail!("unknown syscall number {n}")
				})
			}
		}
	}
}

macro_rules! syscall_bitflags {
	(
		$(
			struct $BitFlags:ident: $T:ty {
				$(
					const $FLAG:ident = $value:expr => $LIBC_FLAG:ident;
				)*
			}
		)*
	) => {
		#[test]
		fn verify_syscall_bitflags() {
			$(
				$BitFlags::verify();
			)*
		}

		$(
			bitflags! {
				struct $BitFlags: $T {
					$(
						const $FLAG = $value;
					)*
				}
			}

			impl $BitFlags {
				#[cfg(test)]
				fn verify() {
					$(
						let libc_flag = stringify!($LIBC_FLAG);
						let flag = stringify!($FLAG);
						assert!(libc_flag.ends_with(flag), "{libc_flag} doesn't end with {flag}");

						let left = Self::$FLAG.bits();
						let right = libc::$LIBC_FLAG;
						assert!(
							left == right,
							"{}::{} ({left:#x}) != libc::{} ({right:#x})", stringify!($BitFlags), stringify!($FLAG), stringify!($LIBC_FLAG),
						);
					)*
				}
			}

			impl SyscallArg for $BitFlags {
				fn try_from_reg(reg: u64) -> Option<Self> {
					SyscallArg::try_from_reg(reg).and_then(Self::from_bits)
				}
			}
		)*
	};
}

trait SyscallArg: Sized {
	fn try_from_reg(reg: u64) -> Option<Self>;
}

impl SyscallArg for u16 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		reg.try_into().ok()
	}
}

impl SyscallArg for u32 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		reg.try_into().ok()
	}
}

impl SyscallArg for u64 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(reg)
	}
}

impl SyscallArg for i32 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(u32::try_from(reg).ok()? as i32)
	}
}

impl SyscallArg for *mut i32 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(usize::try_from_reg(reg)? as *mut i32)
	}
}

impl SyscallArg for usize {
	fn try_from_reg(reg: u64) -> Option<Self> {
		reg.try_into().ok()
	}
}

impl SyscallArg for *const u8 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(usize::try_from_reg(reg)? as *const u8)
	}
}

impl SyscallArg for *mut u8 {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(usize::try_from_reg(reg)? as *mut u8)
	}
}

type FileDesc = i32;

type SigAction = ();
type SysInfo = ();
type Tms = ();
type Stat = ();
type RobustListHead = ();
type RLimit64 = ();

impl SyscallArg for *mut () {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(usize::try_from_reg(reg)? as *mut ())
	}
}

impl SyscallArg for *const () {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(usize::try_from_reg(reg)? as *const ())
	}
}

define_syscalls! {
	enum SyscallEntry {
		fn read(fd: FileDesc, buf: *mut u8, count: usize) -> i64 = 0;
		fn write(fd: FileDesc, buf: *const u8, count: usize) -> i64 = 1;
		fn close(fd: FileDesc) -> i64 = 3;
		fn mmap(addr: u64, len: u64, prot: ProtFlags, flags: MapFlags, fd: FileDesc, off: u64) -> i64 = 9;
		fn mprotect(addr: u64, len: usize, prot: ProtFlags) -> i64 = 10;
		fn brk(brk: u64) -> i64 = 12;
		fn rt_sigaction(sig: i32, act: *const SigAction, oact: *mut SigAction, sigsetsize: usize) -> i64 = 13;
		fn ioctl(fd: FileDesc, cmd: u32, arg: u64) -> i64 = 16;
		fn pread64(fd: FileDesc, buf: *mut u8, count: usize, pos: u64) -> i64 = 17;
		fn access(filename: *const u8, mode: i32) -> i64 = 21;
		fn getcwd(buf: *mut u8, size: u64) -> i64 = 79;
		fn readlink(path: *const u8, buf: *mut u8, bufsiz: i32) -> i64 = 89;
		fn sysinfo(info: *mut SysInfo) -> i64 = 99;
		fn times(tbuf: *mut Tms) -> i64 = 100;
		fn arch_prctl(option: i32, arg2: u64) -> i64 = 158;
		fn set_tid_address(tidptr: *mut i32) -> i64 = 218;
		fn exit_group(error_code: i32) -> i64 = 231;
		fn openat(dfd: FileDesc, filename: *const u8, flags: OpenFlags, mode: u16) -> i64 = 257;
		fn newfstatat(dfd: FileDesc, filename: *const u8, statbuf: *mut Stat, flags: i32) -> i64 = 262;
		fn set_robust_list(head: *mut RobustListHead, len: usize) -> i64 = 273;
		fn prlimit64(pid: i32, resource: u32, new_rlim: *const RLimit64, old_rlim: *mut RLimit64) -> i64 = 302;
		fn getrandom(ubuf: *mut u8, len: usize, flags: GrndFlags) -> i64 = 318;
	}
}

#[derive(Debug, Clone)]
enum EntryExit {
	/// Process is about to enter a syscall
	Entry(SyscallEntry),
	/// Process is about to exit a syscall
	Exit(SyscallEntry, i64),
}

fn main() -> Result<()> {
	let process = Process::spawn(&mut {
		let mut args = env::args();

		// drop argv[0]
		args.next();

		let mut cmd = Command::new(args.next().unwrap());
		for arg in args {
			cmd.arg(arg);
		}

		cmd.env_clear();
		cmd
	})?;

	let options = ptrace::Options::PTRACE_O_TRACESYSGOOD
		| ptrace::Options::PTRACE_O_TRACECLONE
		| ptrace::Options::PTRACE_O_EXITKILL;
	ptrace::setoptions(process.tgid.as_pid(), options)?;

	// this is always equal to tgid for now,
	// but I'm keeping this separate so it's obvious what has to be tgid
	let tid = Tid(process.tgid.0);

	let mut syscall_state: Option<EntryExit> = None;

	loop {
		ptrace::syscall(tid.as_pid(), None)?;
		if let Some(EntryExit::Exit(..)) = syscall_state {
			// syscall has completed now
			syscall_state = None;
		}

		let status = waitpid(tid.as_pid(), Some(WaitPidFlag::__WALL))?;

		match (syscall_state, status) {
			(None, WaitStatus::PtraceSyscall(event_tid)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				let regs = ptrace::getregs(event_tid.as_pid())?;
				let entry = match SyscallEntry::from_regs(regs) {
					Ok(entry) => entry,
					Err(err) => {
						ptrace::kill(event_tid.as_pid())?;
						panic!("{err}");
					}
				};

				if !check_syscall(&process, &entry) {
					ptrace::kill(event_tid.as_pid())?;
					panic!("invalid syscall {entry:?}");
				}

				syscall_state = Some(EntryExit::Entry(entry));
			}
			(Some(EntryExit::Entry(entry)), WaitStatus::PtraceSyscall(event_tid)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				let regs = ptrace::getregs(event_tid.as_pid())?;
				let ret = regs.rax as i64;
				syscall_state = Some(EntryExit::Exit(entry, ret));
			}
			(_, WaitStatus::Exited(event_tid, _)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				// TODO(edef): this only works for main thread
				break;
			}
			(syscall_state, status) => {
				panic!("unknown status {status:?} with syscall_state = {syscall_state:?}")
			}
		}
	}

	Ok(())
}

const AT_FDCWD: i32 = -100;

fn check_syscall(process: &Process, entry: &SyscallEntry) -> bool {
	match *entry {
		SyscallEntry::mmap {
			addr: _,
			len: _,
			prot: _,
			flags,
			fd,
			off: _,
		} => {
			if fd == !0 {
				return flags.contains(MapFlags::ANONYMOUS);
			} else {
				return flags.intersection(MapFlags::PRIVATE | MapFlags::ANONYMOUS)
					== MapFlags::PRIVATE;
			}
		}
		SyscallEntry::ioctl { fd: _, cmd, arg: _ } => {
			match cmd {
				// TCGETS
				0x5401 => {}
				// TIOCGWINSZ
				0x5413 => {}
				_ => return false,
			}
		}
		SyscallEntry::access { filename, mode: _ } => {
			let filename = process.read_mem_cstr(filename as u64).unwrap();
			println!("access({filename:?}, ..)");
		}
		SyscallEntry::readlink {
			path,
			buf: _,
			bufsiz: _,
		} => {
			let path = process.read_mem_cstr(path as u64).unwrap();
			println!("readlink({path:?}, ..)");
		}
		SyscallEntry::arch_prctl { option, arg2: _ } => {
			match option {
				// ARCH_SET_FS
				0x1002 => {}
				_ => return false,
			}
		}
		SyscallEntry::set_tid_address { tidptr: _ } => {
			println!("set_tid_address(..)");
		}
		SyscallEntry::openat {
			dfd,
			filename,
			flags,
			mode: _,
		} => {
			if dfd != AT_FDCWD {
				return false;
			}

			let filename = process.read_mem_cstr(filename as u64).unwrap();

			println!("openat(AT_FDCWD, {filename:?}, {flags:?}, ..)");
		}
		SyscallEntry::newfstatat {
			dfd,
			filename,
			statbuf: _,
			flags: _,
		} => {
			let filename = process.read_mem_cstr(filename as u64).unwrap();
			if dfd == AT_FDCWD {
				println!("newfstatat(AT_FDCWD, {filename:?}, ..)");
			} else if filename.as_bytes() == b"" {
				println!("newfstatat({dfd}, {filename:?})");
			} else {
				return false;
			}
		}
		SyscallEntry::set_robust_list { head: _, len } => {
			if len != 24 {
				panic!("set_robust_list(2) len should be sizeof (struct robust_list_head), actually {len}");
			}
			println!("set_robust_list(..)");
		}
		SyscallEntry::prlimit64 {
			pid,
			resource,
			new_rlim: _,
			old_rlim: _,
		} => {
			if pid != 0 {
				return false;
			}

			match resource as u32 {
				libc::RLIMIT_AS | libc::RLIMIT_STACK | libc::RLIMIT_RSS => {}
				_ => return false,
			}
		}
		SyscallEntry::getrandom {
			ubuf: _,
			len,
			flags,
		} => {
			println!("getrandom(.., {len}, {flags:?})");
		}
		_ => {}
	}
	true
}

syscall_bitflags! {
	struct OpenFlags: i32 {
		const WRONLY  = 1 <<  0 => O_WRONLY;
		const CREAT   = 1 <<  6 => O_CREAT;
		const NOCTTY  = 1 <<  8 => O_NOCTTY;
		const TRUNC   = 1 <<  9 => O_TRUNC;
		const CLOEXEC = 1 << 19 => O_CLOEXEC;
	}

	struct GrndFlags: u32 {
		const NONBLOCK = 1 << 0 => GRND_NONBLOCK;
		const RANDOM   = 1 << 1 => GRND_RANDOM;
	}

	struct MapFlags: i32 {
		const PRIVATE   = 1 <<  1 => MAP_PRIVATE;
		const FIXED     = 1 <<  4 => MAP_FIXED;
		const ANONYMOUS = 1 <<  5 => MAP_ANONYMOUS;
		const DENYWRITE = 1 << 11 => MAP_DENYWRITE;
	}

	struct ProtFlags: i32 {
		const READ  = 1 << 0 => PROT_READ;
		const WRITE = 1 << 1 => PROT_WRITE;
		const EXEC  = 1 << 2 => PROT_EXEC;
	}
}