/* 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 <math.h> /* roundf */
#include <sys/mman.h> /* mmap */
#include <sys/timerfd.h> /* timerfd_create */
#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>

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

#include <ft2build.h>
#include FT_FREETYPE_H

#include "main_font.h"
#include "secondary_font.h"

#include "helpers.h"
#include "arena_allocator.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);
}

int round_float_to_int(float real32)
{
    int result = (int)roundf(real32);
    return result;
}

int str_len(const char *str)
{
    int i = 0;
    while(*(str++) != '\0') {
        i++;
    }
    return i;
}

unsigned long long millis()
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (unsigned long long)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
}

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

void update_crtc_infos(Display *dpy)
{
    if(crtcres) {
        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++) {
            RRCrtc crtc;
            XRRCrtcInfo *info;

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

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);
        }
    }
}

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) {
                }
            }
        }
    }
}

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;
}

/* TODO: Optimize. */
struct fonts {
    FT_Library main_library;
    FT_Face main_face;

    FT_Library secondary_library;
    FT_Face secondary_face;
};

static int load_fonts(struct fonts *fonts)
{
    if(FT_Init_FreeType(&fonts->main_library) != 0)
        return 0;

    if(FT_New_Memory_Face(fonts->main_library,
                          main_ttf, main_ttf_len,
                          0, &fonts->main_face) != 0)
    {
        return 0;
    }

    if(FT_Init_FreeType(&fonts->secondary_library) != 0)
        return 0;

    if(FT_New_Memory_Face(fonts->secondary_library,
                          secondary_ttf, secondary_ttf_len,
                          0, &fonts->secondary_face) != 0)
    {
        return 0;
    }

    return 1;
}

static void clear_background(void *memory, int width, int height, int color)
{
    int x;
    for(x = 0; x < width * height; x++) {
        unsigned int *pixel = &((unsigned int *)memory)[x];
        *pixel = color;
    }
}

static void
draw_glyph(void *memory, int width, int height, int bytes_per_pixel,
           FT_Face face, int x, int y)
{
    int pitch;
    int bitmap_pitch, bitmap_width, bitmap_rows;
    void *bitmap_pixels;

    int row, column;
    int rows, columns;

    unsigned char *tdest;
    int mx, my;

    pitch = width * bytes_per_pixel;

    bitmap_pitch = face->glyph->bitmap.pitch;
    bitmap_width = face->glyph->bitmap.width;
    bitmap_rows = face->glyph->bitmap.rows;
    bitmap_pixels = (void *)face->glyph->bitmap.buffer;
    y -= face->glyph->bitmap_top;

    row = 0;
    column = 0;

    rows = bitmap_rows;
    columns = bitmap_width;

    if(x < 0)
        column = -x;
    if(y < 0)
        row = -y;

    if(x + columns > width)
        columns = width - x;
    if(y + rows > height)
        rows = height - y;
    
    tdest = (unsigned char *)(memory + y*pitch + x*bytes_per_pixel);

    for(my = row; my < rows; my++) {
        unsigned char *source;
        unsigned int *dest;

        source = (unsigned char *)(bitmap_pixels + my*bitmap_pitch);
        dest   = (unsigned int *)(tdest + my*pitch);

        for(mx = column; mx < columns; mx++) {
            unsigned char gray;
            unsigned int source_pixel;

            float a, r, g, b, sa, sr, sg, sb, da, dr, dg, db, ar;

            gray = source[mx];
            source_pixel = (gray << 24) | (gray << 16) | (gray << 8) | gray;

            sa = (float)((source_pixel >> 24) & 0xFF);
            sr = (float)((source_pixel >> 16) & 0xFF);
            sg = (float)((source_pixel >>  8) & 0xFF);
            sb = (float)((source_pixel >>  0) & 0xFF);

            da = (float)((dest[mx] >> 24) & 0xFF);
            dr = (float)((dest[mx] >> 16) & 0xFF);
            dg = (float)((dest[mx] >>  8) & 0xFF);
            db = (float)((dest[mx] >>  0) & 0xFF);

            ar  = sa / 255.0f;

            /* TODO: This should be replaced by premultiplied alpha. */
            a = (1.0f-ar)*da + ar*sa;
            r = (1.0f-ar)*dr + ar*sr;
            g = (1.0f-ar)*dg + ar*sg;
            b = (1.0f-ar)*db + ar*sb;

            dest[mx] = (((unsigned int)round_float_to_int(a) << 24) |
                        ((unsigned int)round_float_to_int(r) << 16) |
                        ((unsigned int)round_float_to_int(g) << 8)  |
                        ((unsigned int)round_float_to_int(b) << 0));
        }
    }
}

static void
draw_string(void *buffer, int width, int height,
            const char *str, FT_Face face, int x, int y)
{
    int str_length, start_x, start_y, i, error;

    str_length = str_len(str);

    start_x = x;
    start_y = y;

    for(i = 0; i < str_length; i++) {
        char character = str[i];

        if(character == ' ') {
            start_x += 20;
        } else {
            int glyph_index = FT_Get_Char_Index(face, character);

            error = FT_Load_Glyph(face, glyph_index, 0);
            if(error)
                break;

            if(face->glyph->format != FT_GLYPH_FORMAT_BITMAP) {
                error = FT_Render_Glyph(face->glyph, FT_RENDER_MODE_NORMAL);
                if(error)
                    break;
            }

            draw_glyph(buffer, width, height, 4, face, start_x, start_y);

            start_x += face->glyph->bitmap.width;
        }
    }
}

static void
get_string_dimension(int *width, int *height, const char *str, FT_Face face)
{
    int str_length, start_x, start_y, i, error;
    int result_width, result_height;

    str_length = str_len(str);

    result_width = 0;
    result_height = 0;

    for(i = 0; i < str_length; i++) {
        char character = str[i];

        if(character == ' ') {
            result_width += 20;
        } else {
            int glyph_index = FT_Get_Char_Index(face, character);

            error = FT_Load_Glyph(face, glyph_index, 0);
            if(error)
                break;

            if(face->glyph->format != FT_GLYPH_FORMAT_BITMAP) {
                error = FT_Render_Glyph(face->glyph, FT_RENDER_MODE_NORMAL);
                if(error)
                    break;
            }

            result_width += face->glyph->bitmap.width;
            if(face->glyph->bitmap.rows > (unsigned int)result_height)
                result_height = face->glyph->bitmap.rows;
        }
    }

    *width = result_width;
    *height = result_height;
}

static void get_time_str(char *buffer, int buffer_size)
{
    time_t now;
    struct tm *tm_info;
    char *time_str;

    now = time(0);
    tm_info = localtime(&now);
    strftime(buffer, buffer_size, "%H:%M", tm_info);
}

/* TODO: We are relying heavily on locks[0], is there a guarantee that it is
 * going to be the main screen? */
static void
render_screen(struct arena *arena,
              Display *dpy, struct lock **locks,
              FT_Face main_font_face, FT_Face secondary_font_face,
              unsigned int pass_len)
{
    XWindowAttributes attrs;
    if(XGetWindowAttributes(dpy, locks[0]->win, &attrs)) {
        int width, height, depth;
        struct arena temp_arena;
        void *data;
        char *time_str;
        int time_str_len;

        width = attrs.width;
        height = attrs.height;
        depth = DefaultDepth(dpy, 0);

        temp_arena = temp_arena_begin(arena);

        Visual *visual = DefaultVisual(dpy, 0);
        GC gc = XCreateGC(dpy, locks[0]->win, 0, NULL);

        data = arena_alloc(&temp_arena, width * height * 4);
        XImage *img = XCreateImage(dpy, visual, depth, ZPixmap, 0,
                                   data, width, height, 32, 0);

        clear_background(data, width, height, 0);

        time_str_len = 64;
        time_str = (char *)arena_alloc(&temp_arena, time_str_len);
        get_time_str(time_str, time_str_len);
        if(FT_Set_Pixel_Sizes(main_font_face, 0, 200) == 0) {
            int mwidth, mheight;
            get_string_dimension(&mwidth, &mheight, time_str, main_font_face);
            draw_string(data,
                        width, height,
                        time_str,
                        main_font_face,
                        width/2-mwidth/2, height/4+mheight/2);
        }

        if(FT_Set_Pixel_Sizes(secondary_font_face, 0, 30) == 0) {
            int mwidth, mheight;

            char *message_str;
            int message_str_len;

            message_str_len = 32;
            message_str = (char *)arena_alloc(&temp_arena, message_str_len);
            memset(message_str, 0, message_str_len);
            if(pass_len > 0) {
                unsigned int i;

                message_str[0] = '\0';
                for(i = 0; i < 20 && i < pass_len; i++)
                    message_str[i] = '*';
                message_str[i] = '\0';
            }

            get_string_dimension(&mwidth, &mheight,
                                 message_str, secondary_font_face);
            draw_string(data,
                        width, height,
                        message_str,
                        secondary_font_face,
                        width/2-mwidth/2, height/2+mheight/2);
        }

        XPutImage(dpy, locks[0]->win, gc, img, 0, 0, 0, 0, width, height);
        XFlush(dpy);

        img->data = 0;
        XDestroyImage(img);
        XFreeGC(dpy, gc);

        temp_arena_end(&temp_arena);
    }
}

/* TODO: Optimize. */
struct dpms_state {
    Bool dpmsavailable;
    CARD16 dpmspowerlevel;
    BOOL dpmsstate;
    CARD16 dpmsstandby, dpmssuspend, dpmsoff;
};

static void dpms_disable(Display *dpy, struct dpms_state *state)
{
    if(DPMSCapable(dpy) &&
       DPMSInfo(dpy, &state->dpmspowerlevel, &state->dpmsstate) &&
       DPMSGetTimeouts(dpy,
                       &state->dpmsstandby,
                       &state->dpmssuspend,
                       &state->dpmsoff))
    {
        state->dpmsavailable = 1;
        if(state->dpmsstate)
            DPMSDisable(dpy);
        XSync(dpy, 0);
    }
}

static void dpms_enable(Display *dpy, struct dpms_state *state)
{
    if(state->dpmsavailable && state->dpmsstate)
    {
        DPMSSetTimeouts(dpy,
                        state->dpmsstandby,
                        state->dpmssuspend,
                        state->dpmsoff);
        DPMSEnable(dpy);
        XSync(dpy, 0);
    }
}

static void update_timerfd(int tfd)
{
    time_t now;
    struct tm tmv;
    int secs_left;
    struct itimerspec its;

    now = time(0);
    localtime_r(&now, &tmv);

    secs_left = 60 - tmv.tm_sec;

    memset(&its, 0, sizeof(its));
    its.it_value.tv_sec = secs_left;
    timerfd_settime(tfd, 0, &its, 0);
}

static void
readpw(struct arena *arena,
       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, tfd, xfd;
    unsigned int len, color;
    struct timeval tv;
    KeySym ksym;
    XEvent ev;
    unsigned long long last_millis;
    struct dpms_state dpms_state;
    struct fonts fonts;
    int error;

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

    xfd = ConnectionNumber(dpy);
    tfd = timerfd_create(CLOCK_REALTIME, 0);

    update_crtc_infos(dpy);

    dpms_disable(dpy, &dpms_state);

    FT_Library main_font_library;
    FT_Face main_font_face;

    FT_Library secondary_font_library;
    FT_Face secondary_font_face;

    if(!load_fonts(&fonts))
        die("slock: load_fonts\n");

    render_screen(arena,
                  dpy, locks,
                  fonts.main_face, fonts.secondary_face,
                  len);
    
    last_millis = millis();
    while(running) {
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(xfd, &fds);
        FD_SET(tfd, &fds);

        XFlush(dpy);

        update_timerfd(tfd);

        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) {
#if DEBUG
                    exit(1);
#else
                    dson = !dson;
                    if(dson)
                        set_screen_on(dpy);
                    else
                        set_screen_off(dpy);
                    continue;
#endif
                }

                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;
                }
            } 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(ev.type == ConfigureNotify) {
            } else {
                for(screen = 0; screen < nscreens; screen++)
                    XRaiseWindow(dpy, locks[screen]->win);
            }

            /* cap redraw at roughly 30 FPS */
            if(millis() - last_millis >= 33) {
                render_screen(arena,
                              dpy, locks,
                              fonts.main_face, fonts.secondary_face,
                              len);
                last_millis = millis();
            }
        }

        int ret;
        if(displaytimeout > 0 && dson) {
            tv.tv_sec = displaytimeout;
            tv.tv_usec = 0;
            ret = select(tfd + 1, &fds, 0, 0, &tv);
        } else {
            ret = select(tfd + 1, &fds, 0, 0, 0);
        }
        if (ret == 0) {
            dson = 0;
            set_screen_off(dpy);
        } else if (ret > 0 && FD_ISSET(xfd, &fds)) {
        } else if (ret > 0 && FD_ISSET(tfd, &fds)) {
            render_screen(arena,
                          dpy, locks,
                          fonts.main_face, fonts.secondary_face,
                          len);
        }
    }

    dpms_enable(dpy, &dpms_state);
}

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);

    /* init */
    wa.override_redirect = 1;
    wa.background_pixel = lock->colors[INIT];
#if DEBUG
    lock->win = XCreateWindow(dpy, lock->root, 0, 0,
                              800,
                              400,
                              0, DefaultDepth(dpy, lock->screen),
                              CopyFromParent,
                              DefaultVisual(dpy, lock->screen),
                              CWBackPixel, &wa);
#else
    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);
#endif
    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);

#if DEBUG
    XMapRaised(dpy, lock->win);
    XSelectInput(dpy, lock->win, KeyPressMask | KeyReleaseMask |
                                 ButtonPressMask | ButtonReleaseMask |
                                 PointerMotionMask | StructureNotifyMask);
    return lock;
#else
    /* 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);
#endif

    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;

    unsigned long long bf_size;
    void *bf;

    bf_size = GB(1);
    bf = mmap(0,
              bf_size,
              PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE,
              -1, 0);
    if(!bf)
        die("slock: mmap\n");

    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 */
    struct arena arena = arena_init(bf, GB(1));
    readpw(&arena, dpy, &rr, locks, nscreens, hash);

    return 0;
}