// 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"))
	}
}

#[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.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 = 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(())
}

const AT_FDCWD: i32 = -100;

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 => {
			let [_fd, command, ..] = entry.args;
			match command {
				// TCGETS
				0x5401 => {}
				// TIOCGWINSZ
				0x5413 => {}
				_ => return false,
			}
		}

		// pread64
		17 => {}

		// access
		21 => {
			let [pathname, _mode, ..] = entry.args;
			let pathname = process.read_mem_cstr(pathname).unwrap();
			println!("access({:?}, ..)", pathname);
		}

		// getcwd
		79 => {}

		// readlink
		89 => {
			let [pathname, _buf, _bufsiz, ..] = entry.args;
			let pathname = process.read_mem_cstr(pathname).unwrap();
			println!("readlink({:?}, ..)", pathname);
		}

		// sysinfo
		99 => {}

		// times
		100 => {}

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

		// set_tid_address
		218 => {
			let [_tidptr, ..] = entry.args;
			println!("set_tid_address(..)");
		}

		// exit_group
		231 => {}

		// openat
		257 => {
			let [dirfd, pathname, flags, _mode, ..] = entry.args;

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

			let pathname = process.read_mem_cstr(pathname).unwrap();

			let flags: i32 = flags.try_into().expect("openat(2) flags don't fit in i32");
			let flags = OpenFlags::from_bits(flags).expect("unknown openat flags");

			println!("openat(AT_FDCWD, {:?}, {:?}, ..)", pathname, flags);
		}

		// newfstatat
		262 => {
			let [dirfd, pathname, _statbuf, _flags, ..] = entry.args;

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

			let pathname = process.read_mem_cstr(pathname).unwrap();

			println!("newfstatat(AT_FDCWD, {:?}, ..)", pathname);
		}

		// set_robust_list
		273 => {
			let [_head, len, ..] = entry.args;
			if len != 24 {
				panic!("set_robust_list(2) len should be sizeof (struct robust_list_head), actually {}", len);
			}
			println!("set_robust_list(..)");
		}

		// prlimit64
		302 => {
			let [pid, resource, _new_limit, _old_limit, ..] = entry.args;

			if pid != 0 {
				return false;
			}

			match resource as u32 {
				libc::RLIMIT_AS | libc::RLIMIT_STACK | libc::RLIMIT_RSS => {}
				_ => return false,
			}
		}

		// getrandom
		318 => {
			let [_buf, buflen, flags, ..] = entry.args;
			let flags = flags
				.try_into()
				.expect("getrandom(2) flags don't fit in u32");
			let flags = GrndFlags::from_bits(flags).expect("unknown getrandom(2) flags");
			println!("getrandom(.., {}, {:?})", buflen, flags);
		}

		_ => return false,
	}
	true
}

bitflags! {
	struct OpenFlags: i32 {
		const WRONLY  = 0o00000001;
		const CREAT   = 0o00000100;
		const NOCTTY  = 0o00000400;
		const TRUNC   = 0o00001000;
		const CLOEXEC = 0o02000000;
	}

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