/* See LICENSE file for license details. */
#define _XOPEN_SOURCE 500
#if HAVE_SHADOW_H
#include <shadow.h>
#endif

#include <ctype.h>
#include <errno.h>
#include <grp.h>
#include <pwd.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <spawn.h>
#include <time.h>
#include <sys/types.h>
#include <X11/extensions/Xrandr.h>
#include <X11/extensions/dpms.h>
#include <X11/keysym.h>
#include <X11/XF86keysym.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>

#include "arg.h"
#include "util.h"

#define LENGTH(X)       (sizeof X / sizeof X[0])
#define CLEANMASK(mask) (mask & ~(numlockmask|LockMask) & (ShiftMask|ControlMask|Mod1Mask|Mod2Mask|Mod3Mask|Mod4Mask|Mod5Mask))

static unsigned int numlockmask = 0;

char *argv0;

enum {
	INIT,
	INPUT,
	FAILED,
	NUMCOLS
};

struct lock {
	int screen;
	Window root, win;
	Pixmap pmap;
	unsigned long colors[NUMCOLS];
};

struct xrandr {
	int active;
	int evbase;
	int errbase;
};

typedef struct {
	unsigned int mod;
	KeySym keysym;
} Passthrough;

typedef union {
    int i;
    unsigned int ui;
    float f;
    const void *v;
} Arg;

typedef struct {
    unsigned int mod;
    KeySym keysym;
    void (*func)(const Arg *);
    const Arg arg;
} Key;

void
spawn(const Arg *arg)
{
    if (fork() == 0) {
        execvp(((char **)arg->v)[0], (char **)arg->v);
    }
}

#include "config.h"

static void
die(const char *errstr, ...)
{
	va_list ap;

	va_start(ap, errstr);
	vfprintf(stderr, errstr, ap);
	va_end(ap);
	exit(1);
}

#ifdef __linux__
#include <fcntl.h>
#include <linux/oom.h>

XRRScreenResources *crtcres = 0;
static int crtccount = 0;
static XRRCrtcInfo *crtcinfos[16];

void update_crtc_infos(Display *dpy)
{
    if(crtcres) {
        /* TODO: Turn on displays here first? */
        int i;
        for(i = 0; i < crtccount; i++) {
            XRRFreeCrtcInfo(crtcinfos[i]);
        }
        memset(crtcinfos, 0, sizeof(XRRCrtcInfo *)*16);
        XRRFreeScreenResources(crtcres);
        crtcres = 0;
        crtccount = 0;
    } else {
        Window root;
        int i;

        root = DefaultRootWindow(dpy);
        crtcres = XRRGetScreenResources(dpy, root);

        for(i = 0; i < crtcres->ncrtc; i++) { /* TODO: FIX: Choose minimum of ncrtc and my crtcinfos value.
                                                 Could it be that the OS disables displays after inactivity? */
            RRCrtc crtc;
            XRRCrtcInfo *info;

            crtc = crtcres->crtcs[i];
            info = XRRGetCrtcInfo(dpy, crtcres, crtc);
            crtcinfos[i] = info;
        }
    }
}

// TODO: Error-prone. See which of these can return errors.
void set_screen_on(Display *dpy)
{
    if(crtcres) {
        int i;
        for(i = 0; i < crtcres->ncrtc; i++) {
            RRCrtc crtc;
            static XRRCrtcInfo *info;

            crtc = crtcres->crtcs[i];
            info = crtcinfos[i];

            XRRSetCrtcConfig(dpy, crtcres, crtcres->crtcs[i], CurrentTime,
                             info->x, info->y,
                             info->mode, info->rotation, info->outputs,
                             info->noutput);
        }
    }
}

// TODO: Error-prone. See which of these can return errors.
void set_screen_off(Display *dpy)
{
    if(crtcres) {
        int i;
        for(i = 0; i < crtcres->ncrtc; i++) {
            RRCrtc crtc;
            static XRRCrtcInfo *info;

            crtc = crtcres->crtcs[i];
            info = crtcinfos[i];

            if(info->mode != None) {
                Status s =
                    XRRSetCrtcConfig(dpy, crtcres, crtcres->crtcs[i],
                                     CurrentTime, 0, 0, None, RR_Rotate_0,
                                     NULL, 0);
                if(s != RRSetConfigSuccess) {
                    /* TODO: Failed to disable... */
                }
            }
        }
    }
}

/* TODO: When new CRTCs are connected when the screen is off. Turn on all
 * screens and then update the structure. We want the displays on when new
 * one connects. */

static void
dontkillme(void)
{
	FILE *f;
	const char oomfile[] = "/proc/self/oom_score_adj";

	if (!(f = fopen(oomfile, "w"))) {
		if (errno == ENOENT)
			return;
		die("slock: fopen %s: %s\n", oomfile, strerror(errno));
	}
	fprintf(f, "%d", OOM_SCORE_ADJ_MIN);
	if (fclose(f)) {
		if (errno == EACCES)
			die("slock: unable to disable OOM killer. "
			    "Make sure to suid or sgid slock.\n");
		else
			die("slock: fclose %s: %s\n", oomfile, strerror(errno));
	}
}
#endif

static const char *
gethash(void)
{
	const char *hash;
	struct passwd *pw;

	/* Check if the current user has a password entry */
	errno = 0;
	if (!(pw = getpwuid(getuid()))) {
		if (errno)
			die("slock: getpwuid: %s\n", strerror(errno));
		else
			die("slock: cannot retrieve password entry\n");
	}
	hash = pw->pw_passwd;

#if HAVE_SHADOW_H
	if (!strcmp(hash, "x")) {
		struct spwd *sp;
		if (!(sp = getspnam(pw->pw_name)))
			die("slock: getspnam: cannot retrieve shadow entry. "
			    "Make sure to suid or sgid slock.\n");
		hash = sp->sp_pwdp;
	}
#else
	if (!strcmp(hash, "*")) {
#ifdef __OpenBSD__
		if (!(pw = getpwuid_shadow(getuid())))
			die("slock: getpwnam_shadow: cannot retrieve shadow entry. "
			    "Make sure to suid or sgid slock.\n");
		hash = pw->pw_passwd;
#else
		die("slock: getpwuid: cannot retrieve shadow entry. "
		    "Make sure to suid or sgid slock.\n");
#endif /* __OpenBSD__ */
	}
#endif /* HAVE_SHADOW_H */

	return hash;
}

static void
readpw(Display *dpy, struct xrandr *rr, struct lock **locks, int nscreens,
       const char *hash)
{
	XRRScreenChangeNotifyEvent *rre;
	char buf[32], passwd[256], *inputhash;
	int num, screen, running, failure, oldc, i, passing, dson, xfd;
	unsigned int len, color;
    struct timeval tv;
	KeySym ksym;
	XEvent ev;

    Bool dpmsavailable;
    CARD16 dpmspowerlevel;
    BOOL dpmsstate;
    CARD16 dpmsstandby, dpmssuspend, dpmsoff;

	len = 0;
	running = 1;
	failure = 0;
	oldc = INIT;
    dson = 1;

    xfd = ConnectionNumber(dpy);

    update_crtc_infos(dpy);

    if(DPMSCapable(dpy) && DPMSInfo(dpy, &dpmspowerlevel, &dpmsstate) &&
       DPMSGetTimeouts(dpy, &dpmsstandby, &dpmssuspend, &dpmsoff))
    {
        dpmsavailable = 1;
        if(dpmsstate)
            DPMSDisable(dpy);
        XSync(dpy, 0);
    }

    /* TODO: Disable the DPMS timeout, save settings, and restore. */

    while(running) {
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(xfd, &fds);

        XFlush(dpy);

        while(XPending(dpy)) {
            XNextEvent(dpy, &ev);

            if(ev.type == MotionNotify) {
                dson = 1;
                set_screen_on(dpy);
            } else if (ev.type == KeyPress) {
                explicit_bzero(&buf, sizeof(buf));
                num = XLookupString(&ev.xkey, buf, sizeof(buf), &ksym, 0);

                if(ksym != XK_Escape) {
                    dson = 1;
                    set_screen_on(dpy);
                }

                if (IsKeypadKey(ksym)) {
                    if (ksym == XK_KP_Enter)
                        ksym = XK_Return;
                    else if (ksym >= XK_KP_0 && ksym <= XK_KP_9)
                        ksym = (ksym - XK_KP_0) + XK_0;
                }
                if (IsFunctionKey(ksym) ||
                    IsKeypadKey(ksym) ||
                    IsMiscFunctionKey(ksym) ||
                    IsPFKey(ksym) ||
                    IsPrivateKeypadKey(ksym))
                    continue;

                passing = 0;
                for (i = 0; i < (int)LENGTH(passthroughs); i++) { /* WARNING: Unsafe cast. */
                    if (ksym == passthroughs[i].keysym &&
                            CLEANMASK(passthroughs[i].mod) == CLEANMASK(ev.xkey.state)) {
                        passing = 1;
                        XSendEvent(dpy, DefaultRootWindow(dpy), True, KeyPressMask, &ev);
                        break;
                    }
                }

                if(passing)
                    continue;

                for (i = 0; i < (int)LENGTH(keys); i++) { /* WARNING: Unsafe cast. */
                    Key key = keys[i];

                    if (ksym == key.keysym &&
                            CLEANMASK(key.mod) == CLEANMASK(ev.xkey.state)) {
                        passing = 1;
                        key.func(&(key.arg));
                    }
                }

                if(passing)
                    continue;

                if(len == 0 && ksym == XK_Escape) {
                    dson = !dson;
                    if(dson)
                        set_screen_on(dpy);
                    else
                        set_screen_off(dpy);
                    continue;
                }

                switch (ksym) {
                case XK_Return:
                    passwd[len] = '\0';
                    errno = 0;
                    if (!(inputhash = crypt(passwd, hash)))
                        fprintf(stderr, "slock: crypt: %s\n", strerror(errno));
                    else
                        running = !!strcmp(inputhash, hash);
                    if (running) {
                        XBell(dpy, 100);
                        failure = 1;
                    }
                    explicit_bzero(&passwd, sizeof(passwd));
                    len = 0;
                    break;
                case XK_Escape:
                    explicit_bzero(&passwd, sizeof(passwd));
                    len = 0;
                    break;
                case XK_BackSpace:
                    if (len)
                        passwd[--len] = '\0';
                    break;
                default:
                    if (num && !iscntrl((unsigned char)buf[0]) &&
                        (len + num < sizeof(passwd))) {
                        memcpy(passwd + len, buf, num);
                        len += num;
                    } else if (buf[0] == '\025') { /* ctrl-u clears input */
                        explicit_bzero(&passwd, sizeof(passwd));
                        len = 0;
                    } else {
                        continue; /* Don't trigger fail screen when pressing control characters */
                    }
                    break;
                }
                color = len ? INPUT : ((failure || failonclear) ? FAILED : INIT);
                if (running && oldc != (int)color) { /* WARNING: Unsafe cast. */
                    for (screen = 0; screen < nscreens; screen++) {
                        XSetWindowBackground(dpy,
                                             locks[screen]->win,
                                             locks[screen]->colors[color]);
                        XClearWindow(dpy, locks[screen]->win);
                    }
                    oldc = color;
                }
		    } else if (rr->active && ev.type == rr->evbase + RRScreenChangeNotify) {
                rre = (XRRScreenChangeNotifyEvent*)&ev;
                for (screen = 0; screen < nscreens; screen++) {
                    if (locks[screen]->win == rre->window) {
                        if (rre->rotation == RR_Rotate_90 ||
                            rre->rotation == RR_Rotate_270)
                        {
                            XResizeWindow(dpy, locks[screen]->win,
                                          rre->height, rre->width);
                        } else {
                            XResizeWindow(dpy, locks[screen]->win,
                                          rre->width, rre->height);
                        }
                        XClearWindow(dpy, locks[screen]->win);
                        break;
                    }
                }
		    } else if(rr->active && ev.type == rr->evbase + RRNotify) {
                /* TODO: Update crts variables and detect if new displays
                 * are connected. XRRNotifyEvent->subtype,
                 * XRRCrtcChangeNotifyEvent, XRROutputChangeNotifyEvent,
                 * RRNotify_OutputProperty*/
            } else {
                for(screen = 0; screen < nscreens; screen++)
                    XRaiseWindow(dpy, locks[screen]->win);
            }
        }

        int ret;
        if(displaytimeout > 0 && dson) {
            tv.tv_sec = displaytimeout;
            tv.tv_usec = 0;
            ret = select(xfd + 1, &fds, 0, 0, &tv);
        } else {
            ret = select(xfd + 1, &fds, 0, 0, 0);
        }
        if (ret == 0) {
            dson = 0;
            set_screen_off(dpy);
        } else if (ret > 0 && FD_ISSET(xfd, &fds)) {
        }
    }

    if(dpmsavailable && dpmsstate)
    {
        DPMSSetTimeouts(dpy, dpmsstandby, dpmssuspend, dpmsoff);
        DPMSEnable(dpy);
        XSync(dpy, 0);
    }
}

static struct lock *
lockscreen(Display *dpy, struct xrandr *rr, int screen)
{
	char curs[] = {0, 0, 0, 0, 0, 0, 0, 0};
	int i, ptgrab, kbgrab;
	struct lock *lock;
	XColor color, dummy;
	XSetWindowAttributes wa;
	Cursor invisible;

	if (dpy == NULL || screen < 0 || !(lock = malloc(sizeof(struct lock))))
		return NULL;

	lock->screen = screen;
	lock->root = RootWindow(dpy, lock->screen);

	for (i = 0; i < NUMCOLS; i++) {
		XAllocNamedColor(dpy, DefaultColormap(dpy, lock->screen),
		                 colorname[i], &color, &dummy);
		lock->colors[i] = color.pixel;
	}

	/* init */
	wa.override_redirect = 1;
	wa.background_pixel = lock->colors[INIT];
	lock->win = XCreateWindow(dpy, lock->root, 0, 0,
	                          DisplayWidth(dpy, lock->screen),
	                          DisplayHeight(dpy, lock->screen),
	                          0, DefaultDepth(dpy, lock->screen),
	                          CopyFromParent,
	                          DefaultVisual(dpy, lock->screen),
	                          CWOverrideRedirect | CWBackPixel, &wa);
	lock->pmap = XCreateBitmapFromData(dpy, lock->win, curs, 8, 8);
	invisible = XCreatePixmapCursor(dpy, lock->pmap, lock->pmap,
	                                &color, &color, 0, 0);
	XDefineCursor(dpy, lock->win, invisible);

    XRRSelectInput(dpy, lock->win, RRScreenChangeNotifyMask |
                                   RRCrtcChangeNotifyMask |
                                   RROutputChangeNotifyMask |
                                   RROutputPropertyNotifyMask);

	/* Try to grab mouse pointer *and* keyboard for 600ms, else fail the lock */
	for (i = 0, ptgrab = kbgrab = -1; i < 6; i++) {
		if (ptgrab != GrabSuccess) {
			ptgrab = XGrabPointer(dpy, lock->root, False,
			                      ButtonPressMask | ButtonReleaseMask |
			                      PointerMotionMask, GrabModeAsync,
			                      GrabModeAsync, None, invisible, CurrentTime);
		}
		if (kbgrab != GrabSuccess) {
			kbgrab = XGrabKeyboard(dpy, lock->root, True,
			                       GrabModeAsync, GrabModeAsync, CurrentTime);
		}

		/* input is grabbed: we can lock the screen */
		if (ptgrab == GrabSuccess && kbgrab == GrabSuccess) {
			XMapRaised(dpy, lock->win);
			if (rr->active)
				XRRSelectInput(dpy, lock->win, RRScreenChangeNotifyMask);

			XSelectInput(dpy, lock->root, SubstructureNotifyMask);
			return lock;
		}

		/* retry on AlreadyGrabbed but fail on other errors */
		if ((ptgrab != AlreadyGrabbed && ptgrab != GrabSuccess) ||
		    (kbgrab != AlreadyGrabbed && kbgrab != GrabSuccess))
			break;

		usleep(100000);
	}

	/* we couldn't grab all input: fail out */
	if (ptgrab != GrabSuccess)
		fprintf(stderr, "slock: unable to grab mouse pointer for screen %d\n",
		        screen);
	if (kbgrab != GrabSuccess)
		fprintf(stderr, "slock: unable to grab keyboard for screen %d\n",
		        screen);
	return NULL;
}

static void
usage(void)
{
	die("usage: slock [-v] [cmd [arg ...]]\n");
}

int
main(int argc, char **argv) {
	struct xrandr rr;
	struct lock **locks;
	struct passwd *pwd;
	struct group *grp;
	uid_t duid;
	gid_t dgid;
	const char *hash;
	Display *dpy;
	int s, nlocks, nscreens;
    int event_base, error_base;

    char *username, *groupname;

	ARGBEGIN {
	case 'v':
		puts("slock-"VERSION);
		return 0;
	default:
		usage();
	} ARGEND

    if(strlen(user) && strlen(group)) {
        username = (char *)user;
        groupname = (char *)group;
    } else {
        username = getenv("USER");
        groupname = username;
    }

	/* validate drop-user and -group */
	errno = 0;
	if (!(pwd = getpwnam(username)))
		die("slock: getpwnam %s: %s\n", username,
		    errno ? strerror(errno) : "user entry not found");
	duid = pwd->pw_uid;
	errno = 0;
	if (!(grp = getgrnam(groupname)))
		die("slock: getgrnam %s: %s\n", groupname,
		    errno ? strerror(errno) : "group entry not found");
	dgid = grp->gr_gid;

#ifdef __linux__
	dontkillme();
#endif

	hash = gethash();
	errno = 0;
	if (!crypt("", hash))
		die("slock: crypt: %s\n", strerror(errno));

	if (!(dpy = XOpenDisplay(NULL)))
		die("slock: cannot open display\n");

    if (!DPMSQueryExtension(dpy, &event_base, &error_base))
		die("slock: dpms extension is not supported\n");

	/* drop privileges */
	if (setgroups(0, NULL) < 0)
		die("slock: setgroups: %s\n", strerror(errno));
	if (setgid(dgid) < 0)
		die("slock: setgid: %s\n", strerror(errno));
	if (setuid(duid) < 0)
		die("slock: setuid: %s\n", strerror(errno));

	/* check for Xrandr support */
	rr.active = XRRQueryExtension(dpy, &rr.evbase, &rr.errbase);

	/* get number of screens in display "dpy" and blank them */
	nscreens = ScreenCount(dpy);
	if (!(locks = calloc(nscreens, sizeof(struct lock *))))
		die("slock: out of memory\n");
	for (nlocks = 0, s = 0; s < nscreens; s++) {
		if ((locks[s] = lockscreen(dpy, &rr, s)) != NULL)
			nlocks++;
		else
			break;
	}
	XSync(dpy, 0);

	/* did we manage to lock everything? */
	if (nlocks != nscreens)
		return 1;

	/* run post-lock command */
	if (argc > 0) {
		pid_t pid;
		extern char **environ;
		int err = posix_spawnp(&pid, argv[0], NULL, NULL, argv, environ);
		if (err) {
			die("slock: failed to execute post-lock command: %s: %s\n",
			    argv[0], strerror(err));
		}
	}

	/* everything is now blank. Wait for the correct password */
	readpw(dpy, &rr, locks, nscreens, hash);

	return 0;
}