// 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 as HostSignal,
			wait::{waitpid, WaitPidFlag, WaitStatus},
		},
		unistd::Pid,
	},
	std::{
		convert::TryInto,
		env,
		ffi::CString,
		fmt::{self, Debug},
		fs::{self, File},
		io::{self, BufRead, Seek, SeekFrom},
		os::unix::process::CommandExt,
		process::Command,
	},
};

mod maps_file;

// 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(_, HostSignal::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"))
	}

	fn read_mappings(&self) -> Result<Vec<maps_file::Mapping>> {
		let pid = self.tgid.as_pid();
		let contents = fs::read_to_string(format!("/proc/{pid}/maps"))?;

		// TODO(edef): consult /proc/$pid/map_files/* for pathnames, since /proc/$pid/maps is unreliable with odd paths
		// we'll want to verify the two against each other, just to be sure they're congruent

		let mut mappings = contents
			.lines()
			.map(maps_file::parse_mapping_line)
			.collect::<Result<_, _>>()?;

		for &mut maps_file::Mapping {
			start,
			end,
			ref mut pathname,
			..
		} in &mut mappings
		{
			if pathname.starts_with('[') && pathname.ends_with(']') {
				// these won't exist in map_files
				continue;
			}

			let map_path = format!("/proc/{pid}/map_files/{start:x}-{end:x}");
			let target = fs::read_link(&map_path)
				.with_context(|| {
					format!("Cannot readlink({map_path:?}) (expected target: {pathname:?})")
				})?
				.into_os_string()
				.into_string()
				.expect("path is not valid UTF-8");
			assert_eq!(*pathname, maps_file::escape_path(&target), "escaping bug?");
			*pathname = target;
		}

		Ok(mappings)
	}

	fn dump_mappings(&self) -> Result<()> {
		let mappings = self.read_mappings()?;
		let mut mappings = mappings.iter().peekable();
		let mut segments = vec![];

		while let Some(mut last) = mappings.next() {
			let mut segment = vec![];
			segment.push(last);
			while let Some(&next) = mappings.peek() {
				if last.dev != next.dev || last.inode != next.inode {
					// not the same file
					break;
				}
				if last.end != next.start || last.offset + last.end - last.start != next.offset {
					// not contiguous
					break;
				}
				last = mappings.next().unwrap();
				segment.push(last);
			}
			segments.push(segment);
		}

		println!("{:#?}", segments);
		Ok(())
	}
}

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(process: &Process, 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 ProcessSyscallArg::try_from_process_reg(process, $arg) {
								Some(x) => x,
								None => bail!("couldn't parse {}(2) {}: {:#08x}", stringify!($syscall), stringify!($arg), $arg)
							}),*
						},
					)*
					(n, _) => bail!("unknown syscall number {n}")
				})
			}
		}
	}
}

#[cfg(test)]
fn libc_check<T: Debug + Eq + fmt::LowerHex>(
	item: &'static str,
	(our_name, our_value): (&'static str, T),
	(libc_name, libc_value): (&'static str, T),
) {
	match () {
		_ if libc_name.ends_with(our_name) => {}
		_ if libc_name.starts_with(&format!("{our_name}_")) => {}
		() => panic!("{libc_name} doesn't match {our_name}"),
	}

	assert!(
		our_value == libc_value,
		"{item}::{our_name} ({our_value:#x}) != libc::{libc_name} ({libc_value:#x})",
	);
}

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() {
					$(
						libc_check(
							stringify!($BitFlags),
							(stringify!($FLAG), Self::$FLAG.bits()),
							(stringify!($LIBC_FLAG), libc::$LIBC_FLAG)
						);
					)*
				}
			}

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

macro_rules! syscall_enums {
	(
		$(
			enum $Enum:ident: $T:ty {
				$(
					$VARIANT:ident = $value:literal $(=> $LIBC_VALUE:ident)?,
				)*
			}
		)*
	) => {
		#[test]
		fn verify_syscall_enums() {
			$(
				$Enum::verify();
			)*
		}

		$(
			#[derive(Debug, Copy, Clone, Eq, PartialEq)]
			#[allow(non_camel_case_types)]
			enum $Enum {
				$($VARIANT = $value),*
			}

			impl SyscallArg for $Enum {
				fn try_from_reg(reg: u64) -> Option<Self> {
					let reg = <$T as SyscallArg>::try_from_reg(reg)?;
					Some(match reg {
						$(
							$value => $Enum::$VARIANT,
						)*
						_ => return None
					})
				}
			}

			impl $Enum {
				#[cfg(test)]
				fn verify() {
					$(
						$(
							libc_check(
								stringify!($Enum),
								(stringify!($VARIANT), Self::$VARIANT as $T),
								(stringify!($LIBC_VALUE), libc::$LIBC_VALUE)
							);
						)?
					)*
				}
			}
		)*
	};
}

trait ProcessSyscallArg: Sized {
	fn try_from_process_reg(process: &Process, reg: u64) -> Option<Self>;
}

impl ProcessSyscallArg for CString {
	fn try_from_process_reg(process: &Process, reg: u64) -> Option<Self> {
		process.read_mem_cstr(reg).ok()
	}
}

impl<T: SyscallArg> ProcessSyscallArg for T {
	fn try_from_process_reg(_process: &Process, reg: u64) -> Option<Self> {
		SyscallArg::try_from_reg(reg)
	}
}

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 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: Option<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: Signal, act: *const SigAction, oact: *mut SigAction, sigsetsize: usize) -> i64 = 13;
		fn ioctl(fd: FileDesc, cmd: Ioctl, arg: u64) -> i64 = 16;
		fn pread64(fd: FileDesc, buf: *mut u8, count: usize, pos: u64) -> i64 = 17;
		fn access(filename: CString, mode: AccessMode) -> i64 = 21;
		fn getcwd(buf: *mut u8, size: u64) -> i64 = 79;
		fn readlink(path: CString, 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: ArchOption, arg2: u64) -> i64 = 158;
		fn set_tid_address(tidptr: *mut i32) -> i64 = 218;
		fn exit_group(error_code: i32) -> i64 = 231;
		fn openat(dfd: DirFd, filename: CString, flags: OpenFlags, mode: FileMode) -> i64 = 257;
		fn newfstatat(dfd: DirFd, filename: CString, statbuf: *mut Stat, flags: AtFlags) -> i64 = 262;
		fn set_robust_list(head: *mut RobustListHead, len: usize) -> i64 = 273;
		fn prlimit64(pid: i32, resource: ResourceLimit, 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
	})?;

	process.dump_mappings()?;

	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(&process, regs) {
					Ok(entry) => entry,
					Err(err) => {
						ptrace::kill(event_tid.as_pid())?;
						panic!("{err}");
					}
				};

				if !check_syscall(&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(entry: &SyscallEntry) -> bool {
	match *entry {
		SyscallEntry::mmap {
			addr,
			len: _,
			prot: _,
			flags,
			fd,
			off: _,
		} => {
			if addr % 4096 != 0 {
				return false;
			}

			match fd {
				None => {
					return flags.contains(MapFlags::ANONYMOUS);
				}
				Some(_) => {
					return flags.intersection(MapFlags::PRIVATE | MapFlags::ANONYMOUS)
						== MapFlags::PRIVATE;
				}
			}
		}
		SyscallEntry::mprotect { addr, len, prot: _ } => {
			return addr % 4096 == 0 && len % 4096 == 0;
		}
		SyscallEntry::rt_sigaction {
			sig: _,
			act: _,
			oact: _,
			sigsetsize,
		} => {
			if sigsetsize != 8 {
				panic!(
					"rt_sigaction(2) sigsetsize should be sizeof (sigset_t), actually {sigsetsize}"
				);
			}
		}
		SyscallEntry::access { ref filename, mode } => {
			println!("access({filename:?}, {mode:?})");
		}
		SyscallEntry::readlink {
			ref path,
			buf: _,
			bufsiz: _,
		} => {
			println!("readlink({path:?}, ..)");
		}
		SyscallEntry::set_tid_address { tidptr: _ } => {
			println!("set_tid_address(..)");
		}
		SyscallEntry::openat {
			dfd,
			ref filename,
			flags,
			mode: _,
		} => {
			if dfd != DirFd::Cwd {
				return false;
			}

			println!("openat({dfd:?}, {filename:?}, {flags:?}, ..)");
		}
		SyscallEntry::newfstatat {
			dfd,
			ref filename,
			statbuf: _,
			flags,
		} => {
			match (dfd, filename.as_bytes()) {
				(_, b"") if !flags.contains(AtFlags::EMPTY_PATH) => {
					// empty path without AT_EMPTY_PATH
					return false;
				}
				(DirFd::Cwd, _) | (_, b"") => {
					println!("newfstatat({dfd:?}, {filename:?}, .., {flags:?})");
				}
				_ => {
					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;
			}
		}
		SyscallEntry::getrandom {
			ubuf: _,
			len,
			flags,
		} => {
			println!("getrandom(.., {len}, {flags:?})");
		}
		SyscallEntry::times { .. } => {}
		_ => println!("{entry:?}"),
	}
	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;
	}

	struct AtFlags: i32 {
		const EMPTY_PATH = 1 << 12 => AT_EMPTY_PATH;
	}

	struct AccessMode: i32 {
		const F = 0 => F_OK;
		const X = 1 << 0 => X_OK;
		const W = 1 << 1 => W_OK;
		const R = 1 << 2 => R_OK;
	}

	struct FileMode: u32 {
		const IRUSR = 0o400 => S_IRUSR;
		const IWUSR = 0o200 => S_IWUSR;
		const IXUSR = 0o100 => S_IXUSR;

		const IRGRP = 0o040 => S_IRGRP;
		const IWGRP = 0o020 => S_IWGRP;
		const IXGRP = 0o010 => S_IXGRP;

		const IROTH = 0o004 => S_IROTH;
		const IWOTH = 0o002 => S_IWOTH;
		const IXOTH = 0o001 => S_IXOTH;
	}
}

syscall_enums! {
	enum ResourceLimit: u32 {
		STACK = 0x3 => RLIMIT_STACK,
		RSS   = 0x5 => RLIMIT_RSS,
		AS    = 0x9 => RLIMIT_AS,
	}

	enum ArchOption: i32 {
		SET_FS = 0x1002,
	}

	enum Ioctl: u64 {
		TCGETS     = 0x5401 => TCGETS,
		TIOCGWINSZ = 0x5413 => TIOCGWINSZ,
	}

	enum Signal: i32 {
		ILL  =  4 => SIGILL,
		ABRT =  6 => SIGABRT,
		BUS  =  7 => SIGBUS,
		FPE  =  8 => SIGFPE,
		SEGV = 11 => SIGSEGV,
	}
}

#[derive(Clone, Copy, Eq, PartialEq)]
struct FileDesc(i32);

impl Debug for FileDesc {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "{}", self.0)
	}
}

impl SyscallArg for FileDesc {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(match i32::try_from_reg(reg)? {
			fd @ 0..=i32::MAX => FileDesc(fd),
			_ => return None,
		})
	}
}

impl SyscallArg for Option<FileDesc> {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(match i32::try_from_reg(reg)? {
			-1 => None,
			fd @ 0..=i32::MAX => Some(FileDesc(fd)),
			_ => return None,
		})
	}
}

#[derive(Clone, Copy, Eq, PartialEq)]
enum DirFd {
	Cwd,
	Fd(FileDesc),
}

impl Debug for DirFd {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			DirFd::Cwd => write!(f, "AT_FDCWD"),
			DirFd::Fd(FileDesc(fd)) => write!(f, "{fd}"),
		}
	}
}

impl SyscallArg for DirFd {
	fn try_from_reg(reg: u64) -> Option<Self> {
		Some(match i32::try_from_reg(reg)? {
			AT_FDCWD => Self::Cwd,
			fd @ 0..=i32::MAX => DirFd::Fd(FileDesc(fd)),
			_ => return None,
		})
	}
}