/* [ʞ] xpriv.c <c.hale.su/lexi/util>
 *  ~ lexi hale <lexi@hale.su>
 *  $ cc -Ofast xpriv.c -lrt -lutil -lX11 -oxpriv
 *  © affero general public license

 * xpriv.c is a tool for a very specific use case i have.
 * for security's  sake, i  don't tie my  ssh keys  to my
 * login session.  when i intend  to use them, i  spawn a
 * special terminal  and load the keys  into memory, then
 * exit that terminal when  i'm done using them. however,
 * i often lose  track of or accidentally  kill that win-
 * dow,  despite  adding  a  visual cue  to  its  prompt.
 * safekill.c was one half of  my attempt to address this
 * problem; this is the other.
 *
 * xpriv performs several different tasks to accomplish a
 * single  purpose: i  can  hit a  single keystroke  that
 * will either  conjure up  a new privileged  session, or
 * switch to one  that's already active if  it exists. it
 * does this  by first  checking for  the existence  of a
 * shared  memory  segment. if  it  doesn't  find it,  it
 * starts  a  new  session;  if it  *does*  find  it,  it
 * retrieves the  X11 window  ID from that  shared memory
 * and sends  a _NET_ACTIVE_WINDOW client message  to the
 * root  X  window.  the window  manager  interprets  the
 * message, activating the window.
 *
 * the  flag -k  can also  be passed,  in which  case the
 * utility instructs the running process to liquidate its
 * subprocesses and exit itself.
 *
 * if the shared  memory does not exist,  xpriv creates a
 * new instance  of urxvt. this  instance is told  to run
 * the command  “xpriv -a”  instead of the  user’s normal
 * shell. the -a flag instructs xpriv to get the terminal
 * window’s  ID from  the $WINDOWID  environment variable
 * which urxvt  sets. after this, a  ssh-agent process is
 * launched. xpriv waits until it has opened a socket and
 * then runs ssh-add without parameters to add the user's
 * default keys to the session.
 *
 * after  a success  key-add  has  been confirmed,  xpriv
 * marks the window as “vital”  by setting the X property
 * “_k_vital” on the  window. if the login  fails or does
 * not complete,  safekill.c will  still terminate  it at
 * any time.  the vital  flag is removed  as soon  as the
 * controlling shell terminates; it does *not* remain for
 * the lifetime  of the window, so  a "temporary" session
 * can  be created  in  the current  terminal by  calling
 * `xpriv -a` directly.
 *
 * xpriv  does  its  best   to  clean  up  after  itself,
 * killing all sensitive processes and their children and
 * removing the  shmem segment  when it  is no  longer in
 * use, even if exits somewhat abnormally. if you have to
 * kill -9 xpriv  at any point tho, you can  make it work
 * again (on linux) with
 *   rm /dev/shm/k.xpriv:(xpriv binary basename)
 *
 * you can have multiple xpriv sessions by creating soft-
 * links to the binary with a different name for each.

 TODO send signal to urxvtd directly instead of launching
      urxvtc as a separate process
	  
 TODO make shell & commands performed configurable, with
      flags to control or supplant ssh-agent 

 TODO add a randomizer call that works on BSD
 
 TODO document flags
 
 TODO implement/remove lock flag

 TODO add flag to bring window to current desktop

 TODO rewrite using sysv shmem */
 

#include <pwd.h>
#include <pty.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/shm.h>
#include <X11/Xlib.h>
#include <errno.h>
#include <sys/random.h>

#ifdef _RAND_SYSCALL
#	include <sys/syscall.h>
#	define getrandom(a,b,c) syscall(SYS_getrandom,a,b,c)
	/* this is necessary on certain platforms due to
	 * certainly ungodly libc issues, i think. */
#endif

#ifdef _SHM_LINUX
	/* xpriv was originally written with linux shared
	 * memory in mind, because i'm an idiot. i've since
	 * redesigned it to use the superior old sysv shm
	 * api that works on every other unix, not just
	 * linux, but if for some ungodly reason you want
	 * to use the linux one, then just pass -D_SHM_LINUX */
#	define shmem_prefix "/k.xpriv:"
#	define shmlinux(...) __VA_ARGS__
#	define shmsysv(...)
#else
#	define shmlinux(...) 
#	define shmsysv(...) __VA_ARGS__
#endif

enum /* constants */ {
	false = 0, true = 1,
	shmsysv(shmem_key = 0x53373EC3,) /* ye olde magique numbre */
};

typedef _Bool bool;
enum mode { mode_usage, mode_go, mode_register,
	mode_kill, mode_lock };

enum res { ok, fail_parse, fail_arg, fail_opt, fail_shm,
               fail_mem, fail_pty, fail_nop, fail_win,
               fail_wid, fail_x11};
enum res bad(enum res code) {
	if (code == ok) return ok;

	write(1,"\e[1m|e|\e[0m ",12);
	const char* msg; uint8_t len;
	switch(code) {
		#define say(x) { msg = (x "\n"); len = (sizeof x) + 1; break; }
		case fail_parse: say("could not parse command-line options; pass -u for usage");
		case fail_arg  : say("invalid argument provided");
		case fail_opt  : say("no such option");
		case fail_shm  : say("instance already running");
		case fail_mem  : say("could not map memory");
		case fail_pty  : say("could not alloc pty");
		case fail_nop  : say("no instance running");
		case fail_win  : say("cannot open display");
		case fail_wid  : say("WINDOWID not defined");
		case fail_x11  : say("X server refuses to handle request");
		#undef say
	}
	write(1, msg, len);
	return code;
}

struct signal {
	Window    wid;
	pid_t     pid;
	pid_t     agent;
	enum mode op;
}*global;

char* itoa(unsigned long i, char* buf, size_t bufsz) {
	char* cur = buf + bufsz;
	while (cur >= buf) {
		*--cur='0' + i % 10;
		if (i < 10) break; else i /= 10;
	}
	return cur;
}

bool run;

struct termios initial_state;

void sigusr(int a) { if (global -> op == mode_kill) run = false; }
void sigterm(int a) { run = false; }

void spawn(pid_t ssha, const char* const sockn) {
	if (ssha = fork()) {
		char pid_s_buf[16];
		char* pid_s = itoa(ssha, pid_s_buf, sizeof(pid_s_buf));
		while (access(sockn, F_OK)); // avoid nasty race condition
		setenv("SSH_AGENT_PID", pid_s, true);
		setenv("SSH_AUTH_SOCK", sockn, true);
		global -> agent = ssha;
	} else {
		close(1); close(0);
		execlp("ssh-agent","ssh-agent","-D","-a",sockn,0);
	}
}

enum res register_window(shmlinux(const char* const id,) bool weak, const char* name) {
	// the id field denotes the path to the shared memory
	// in use, and allows the user to have multiple
	// contexts by creating aliases to the binary
	struct signal* s;
	shmlinux ({
		int fd = shm_open(id, O_CREAT | O_EXCL | O_RDWR, 0600);
		ftruncate(fd, sizeof(struct signal));
		if (fd == -1) return fail_shm;
		
		s = mmap(0, sizeof(struct signal),
				PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
		if(s == MAP_FAILED) return fail_mem;
		else global = s;
	}) shmsysv (int shmem;
		shmem = shmget(shmem_key, sizeof (struct signal), IPC_CREAT | 0777);
		if (shmem == -1) return fail_shm;

		s = shmat(shmem, 0,0);

		if (s == (void*)-1) return fail_mem;
		else global = s;
	)

	Display* xdpy = XOpenDisplay(NULL);
	Atom xvital;

	/* x11 */ {
		if (xdpy == NULL) return fail_win;
		Window win;
		const char* winid_s;
		if (!(winid_s = getenv("WINDOWID"))) return fail_wid;
		win=(Window)strtol(winid_s,NULL,10);
		if(weak == false) xvital = XInternAtom(xdpy,"_k_vital",false);
		s -> wid = win;
	}

	pid_t child;
	if(child = fork()) {
		s -> pid = child;
		sigset_t mask, oldmask;
		sigemptyset(&mask);
		sigaddset(&mask, SIGUSR1);
		sigprocmask(SIG_BLOCK, &mask, &oldmask);

		signal(SIGUSR1, sigusr);
		signal(SIGTERM, sigterm);
		signal(SIGCHLD, sigterm);
		signal(SIGHUP, sigterm);

		run = true;
		while(run == true) {
			sigsuspend(&oldmask);
		}

		if(s -> op == mode_kill) {
			kill(s -> pid, SIGTERM);
		}
		kill(s -> agent, SIGTERM);
		sigprocmask(SIG_BLOCK, &mask, NULL);

		shmlinux(shm_unlink(id));
		shmsysv({
			struct shmid_ds bleh; /* ?? */
			shmctl(shmem, IPC_RMID, &bleh);
		})

		if (!weak) {
			XDeleteProperty(xdpy,s -> wid,xvital);
			XSync(xdpy,false);
		}
		XCloseDisplay(xdpy);
	} else {
		// now we start ssh-agent and set the proper environment
		// variables
		pid_t ssha;

		if (name == NULL) {
			/* messy part */ 
			const char* tmp; //tmpdir defined?
			if (!(tmp = getenv("XDG_RUNTIME_DIR")))
				if (!(tmp = getenv("TMPDIR"))) tmp = "/tmp";

			size_t tmpsz = strlen(tmp);
			
			char* sockn = malloc(tmpsz + 1 + sizeof "ssh." + 11);
			strcpy(sockn, tmp);
			sockn[tmpsz] = '/';
			strcpy(sockn+tmpsz+1,"ssh.");
			
			// first, gen a random identifier so we have the ability
			// to know where the socket winds up
			uint8_t* rndid = sockn+tmpsz+1+4;
			getrandom(rndid, 11, 0);
			
			// assuming ascii…
			for(uint8_t*i=rndid;i<rndid+11;++i) {
				*i = '0' + (*i % (25 * 2 + 10));
				if (*i > '9') *i += 7;
				if (*i > 'Z') *i += ('a' - 'Z');
			}

			name = sockn;
		}
		spawn(ssha, name);
	
		pid_t sad;
		int p;
		if (sad = fork()) {
			int added;
			waitpid(sad, &added, 0);
			if (added == 0) {
				if (weak == false) {
					XChangeProperty(xdpy,s -> wid,xvital,xvital,8,PropModeReplace,"\01", 1);
					XSync(xdpy,false);
				}
				write(1,"\033c",3);
				execlp("fish","fish",NULL);
			}
		} else {
			execlp("ssh-add","ssh-add",NULL);
		}
	}

	return ok;
}
enum res kill_window(shmlinux(const char* id)) {
	struct signal* s;
	shmlinux({
		int fd = shm_open(id, O_RDWR, 0600);
		if (fd == -1) return fail_nop;

		s = mmap(0, sizeof(struct signal),
				PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
		if(s == MAP_FAILED) return fail_mem;
	}) shmsysv (int shmem; {
		shmem = shmget(shmem_key, sizeof(struct signal), 0);
		if (shmem == -1) return fail_nop;
		s = shmat(shmem,0,0);
		if (s == (void*)-1) return fail_mem;
	})

	shmsysv({
		struct shmid_ds bleh; /* ?? */
		shmctl(shmem, IPC_RMID, &bleh);
	})

	s -> op = mode_kill;
	kill(s -> pid, SIGUSR1);

	return ok;
}
enum res activate_window(Window w) {
	Display* dpy = XOpenDisplay(NULL);
	long mask = SubstructureRedirectMask | SubstructureNotifyMask;
	XEvent ev = {
		.xclient.type       = ClientMessage,
		.xclient.serial     = 0,
		.xclient.send_event = True,
		.xclient.message_type = XInternAtom(dpy, "_NET_ACTIVE_WINDOW", False),
		.xclient.display    = dpy,
		.xclient.window     = w,
		.xclient.format     = 32,
		.xclient.data.l[0]  = 2,
		.xclient.data.l[1]  = CurrentTime,
		.xclient.data.l[2]  = 0,
	};
	
	if(!XSendEvent(dpy, DefaultRootWindow(dpy), False, mask, &ev)) return fail_x11;

	XSync(dpy,false);
	return ok;
}

int main(int sz, char** argv) {
	char* binname = "xpriv";
	if(sz > 0) binname = argv[0];

	enum mode op = mode_go;
	bool init_weak = false;
	const char* init_named = NULL;

	for(int i = 1; i<sz; ++i) {
		char* v = argv[i];
		if (*v != '-') return bad(fail_arg);
		char* opt = v + 1;
		bool seen_string_arg = false;
	read_opt:
		switch(*opt) {
			case 'a': op = mode_register; break;
			case 'k': op = mode_kill; break;
			case 'l': op = mode_lock; break;
			case 'h': op = mode_usage; break;
			case 'w': init_weak = true; break;
			case 'n': if (i == sz -1 || seen_string_arg)
						  return bad(fail_arg);
					  init_named = argv[i+1];
					  seen_string_arg = true;
					  ++ i; break;
			default: return bad(fail_opt);
		}
		if(opt[1] != 0) { ++opt; goto read_opt; }
	}

	shmlinux(
		size_t nsz;
		const char* basename = argv[0], *p;
		for (p = binname; *p!=0; ++p) {
			if(*p == '/') basename = p + 1;
		}
		nsz = p - basename;
		char shid[nsz + sizeof shmem_prefix];
		strncpy(shid,shmem_prefix,sizeof shmem_prefix);
		strncpy(shid + sizeof shmem_prefix - 1, basename, nsz);
		shid[nsz + sizeof shmem_prefix - 1] = 0;
	) shmsysv ({
		/* TODO: implement ability to have multiple
		 * keys based on the name of the binary */
	})

	tryagain: if (op == mode_go) {
		int shm;
		if ((shm = shmlinux(
				shm_open(shid, O_RDWR, 0600)
			) shmsysv (
				shmget(shmem_key,sizeof (struct signal),0)
			)) == -1) {
			const char* args[] = {
				"urxvtc", "-bg", "[90]#4b0024",
				          "-e",  argv[0],
						  (init_weak?"-aw":"-a"), 0, 0, 0};

			const uint8_t argsz = sizeof args/sizeof(const char*);
			if (init_named != NULL) {
				args[argsz - 3] = "-n"; // im sorry
				args[argsz - 2] = init_named;
			}

			execvp("urxvtc", (char* const*)args);
		} else {
			struct shmid_ds stat;
			shmctl(shm, IPC_STAT, &stat);
			if (stat.shm_nattch == 0) {
				/* if the shm segment is not attached to any
				 * process, it's a relic that needs to be
				 * cleaned up before we do anything else. on
				 * a sane OS, there would be built-in non-
				 * persistence mechanisms for shared mem, but
				 * alas, as you already know, POSIX */
				shmctl(shm, IPC_RMID, &stat);
				goto tryagain;
			}

			struct signal*s = shmlinux(
				mmap(0, sizeof(struct signal), PROT_READ | PROT_WRITE, MAP_SHARED, shm, 0)
			) shmsysv (
				shmat(shm, 0, 0)
			);

			if (s == shmlinux(MAP_FAILED) shmsysv((void*)-1))
				return bad(fail_mem);
			return bad(activate_window(s->wid));
		}
	} else if (op == mode_register)
		return bad(register_window(shmlinux(shid,) init_weak,init_named));
	else if (op == mode_kill)
		return bad(kill_window(shmlinux(shid)));
	else if (op == mode_usage) {
		write(1,"\e[1musage:\e[0m ",15);
		write(1, argv[0], strlen(argv[0]));
		write(1, " [-aklw [arg]]\n",16);
		return fail_parse;
	}
}