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

use {
	anyhow::{bail, Context, Result},
	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)?;

		// TODO(V): replace with CString::from_vec_with_nul when we update to Rust 1.58
		if buf.pop() != Some(0) {
			bail!("Couldn't find null terminator");
		}

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

#[derive(Debug, Copy, Clone)]
struct SyscallEntry {
	number: u64,
	// rdi, rsi, rdx, rcx, r8, r9
	args: [u64; 6],
}

impl SyscallEntry {
	fn from_regs(regs: libc::user_regs_struct) -> SyscallEntry {
		SyscallEntry {
			number: regs.orig_rax,
			args: [regs.rdi, regs.rsi, regs.rdx, regs.rcx, regs.r8, regs.r9],
		}
	}
}

#[derive(Debug, Copy, 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
	})?;

	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 = SyscallEntry::from_regs(regs);

				syscall_state = Some(EntryExit::Entry(entry));

				if !check_syscall(&process, entry) {
					ptrace::kill(event_tid.as_pid())?;
					panic!("unsupported syscall {:?}", 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;
			}
			_ => panic!(
				"unknown status {:?} with syscall_state = {:?}",
				status, syscall_state
			),
		}
	}

	Ok(())
}

fn check_syscall(process: &Process, entry: SyscallEntry) -> bool {
	match entry.number {
		// read
		0 => {}

		// write
		1 => {}

		// close
		3 => {}

		// mmap
		9 => {
			let [_addr, _len, _prot, flags, fd, _off] = entry.args;
			if fd != !0 {
				return flags & (libc::MAP_PRIVATE as u64) != 0;
			} else {
				return flags & (libc::MAP_ANON as u64) != 0;
			}
		}

		// mprotect
		10 => {}

		// brk
		12 => {}

		// rt_sigaction
		13 => {}

		// ioctl
		16 => match entry.args[1] {
			// TCGETS
			0x5401 | 0x5413 => {}
			_ => return false,
		},

		// pread64
		17 => {}

		// access
		21 => {}

		// getcwd
		79 => {}

		// readlink
		89 => {}

		// sysinfo
		99 => {}

		// times
		100 => {}

		// arch_prctl
		158 => match entry.args[0] {
			// ARCH_SET_FS
			0x1002 => {}
			_ => return false,
		},

		// exit_group
		231 => {}

		// openat
		257 => {
			let [dirfd, pathname, _flags, _mode, _, _] = entry.args;
			const AT_FDCWD: i32 = -100;

			if dirfd.try_into() == Ok(AT_FDCWD) {
				return false;
			}

			println!(
				"openat(AT_FDCWD, {:?}, ..)",
				process.read_mem_cstr(pathname).unwrap()
			);
		}

		// newfstatat
		262 => {}

		// prlimit64
		302 => {
			// pid = 0
			if entry.args[0] != 0 {
				return false;
			}

			match entry.args[1] as u32 {
				libc::RLIMIT_AS | libc::RLIMIT_STACK | libc::RLIMIT_RSS => {}
				_ => return false,
			}
		}

		_ => return false,
	}
	true
}