/* [ʞ] 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 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. spriv waits until it has opened a socket and
 * then runs ssh-add without parameters to add the user's
 * default keys to the session.
 *
 * 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 */
 

#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 <X11/Xlib.h>
#include <errno.h>
#include <sys/random.h> //TODO bsd compat

#define shmem_prefix "/k.xpriv:"

typedef enum { false, true } 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; }

enum res register_window(const char* id, bool weak) {
	// 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
	int fd = shm_open(id, O_CREAT | O_EXCL | O_RDWR, 0600);
	ftruncate(fd, sizeof(struct signal));
	if (fd == -1) return fail_shm;
	
	struct signal* s = mmap(0, sizeof(struct signal),
			PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if(s == MAP_FAILED) 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);
		shm_unlink(id);
		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;

		/* messy part */ {
			const char* tmp; //tmpdir defined?
			if (!(tmp = getenv("TMPDIR"))) tmp = "/tmp";
			size_t tmpsz = strlen(tmp);
			
			char sockn[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');
			}

			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);
				s -> agent = ssha;
			} else {
				close(1); close(0);
				execlp("ssh-agent","ssh-agent","-D", "-a",sockn,0);
			}
		}
		
		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 { return ok; }
		} else {
			execlp("ssh-add","ssh-add",NULL);
		}
	}

	return ok; // this is kind of pointless but w/e
}
enum res kill_window(const char* id) {
	int fd = shm_open(id, O_RDWR, 0600);
	if (fd == -1) return fail_nop;

	struct signal* s = mmap(0, sizeof(struct signal),
			PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if(s == MAP_FAILED) return fail_mem;

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

	return ok;
}
enum res activate_window(Window w) {
	Display* dpy = XOpenDisplay(NULL);
	XEvent ev;
	long mask = SubstructureRedirectMask | SubstructureNotifyMask;
	ev.xclient.type       = ClientMessage;
	ev.xclient.serial     = 0;
	ev.xclient.send_event = True;
	ev.xclient.message_type = XInternAtom(dpy, "_NET_ACTIVE_WINDOW", False);
	ev.xclient.display    = dpy;
	ev.xclient.window     = w;
	ev.xclient.format     = 32;
	ev.xclient.data.l[0]  = 2;
	ev.xclient.data.l[1]  = CurrentTime;
	ev.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) {
	enum mode op = mode_go;
	bool init_weak = false;

	for(int i = 1; i<sz; ++i) {
		char* v = argv[i];
		if (*v != '-') return bad(fail_arg);
		char* opt = v + 1;
	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;
			default: return bad(fail_opt);
		}
		if(opt[1] != 0) { ++opt; goto read_opt; }
	}

	size_t nsz;
	const char* basename, *p;
	for (p = argv[0]; *p!=0; ++p) {
		if(*p == '/') basename = p + 1;
	}
	nsz = p - basename;
	char shid[nsz + sizeof shmem_prefix + 0];
	strncpy(shid,shmem_prefix,sizeof shmem_prefix);
	strncpy(shid + sizeof shmem_prefix - 1, basename, nsz);

	if (op == mode_go) {
		int fd;
		if ((fd = shm_open(shid, O_RDWR, 0600)) == -1) {
			execlp("urxvtc", "urxvtc", "-bg", "[80]#4b0024",
					"-e", argv[0], (init_weak?"-aw":"-a"), 0);
		} else {
			struct signal*s = mmap(0, sizeof(struct signal),
					PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
			return bad(activate_window(s->wid));
		}
	} else if (op == mode_register)
		return bad(register_window(shid,init_weak));
	else if (op == mode_kill)
		return bad(kill_window(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;
	}
}

