#include "Runner.h"
#include "IsolationManager.h"

#include <unistd.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <signal.h>

#include <QCommandLineParser>
#include <QApplication>

// #define DEBUG_MESSAGE

static void sigchld_hdl (int sig)
{
    // Wait for all dead processes.
    // We use a non-blocking call to be sure this signal handler will never block.
    while (waitpid(-1, NULL, WNOHANG) > 0) { }
}

static FILE* findJail(std::deque<std::string> &jailCells, const std::string &pipeFilename)
{
    for (auto i = jailCells.begin(); i != jailCells.end(); ++i) {
        if (*i == pipeFilename) {
            struct stat fileData;
            if (lstat(pipeFilename.c_str(), &fileData)) {
                jailCells.erase(i);
                return nullptr;
            }
            return fopen(pipeFilename.c_str(), "w");
        }
    }
    return nullptr;
}

std::string pipeFilename(uint32_t ownerUid, const Message &message)
{
    // create the pipe-file
    for (auto i = message.iBegin(); i.isValid(); i.next()) {
        if (i.isJailId()) {
            uint32_t jailId = i.jailId();
            char path[50];
            auto home = getenv("HOME");
            assert(home);
            snprintf(path, sizeof(path), "%s/.local/jails/%u/.iso-pipe", home, jailId);
            return std::string(path);
        }
    }
    return std::string();
}

/*
 The runner is the priviledged part that starts processes in their own
 settings.
*/
static void mainRunner(char **argv, int inputId, int outputId)
{
    // check basic correctness.
    if (getenv("HOME") == nullptr) {
        fprintf(stderr, "Runner: missing HOME env var\n");
        return;
    }

    const uint32_t ownerUid = getuid(); // the real user, not root.

    std::deque<std::string> occupiedJailcells;

    /* Create a handler that takes all forked-off processes and gets
       notified when they finish after which we call the sigchld_handler
       method.
       This method will 'wait' on all children, which is required for the
       kernel to cleanup after the children after they finish.
     */
    struct sigaction act;
    memset (&act, 0, sizeof(act));
    act.sa_handler = sigchld_hdl;
    if (sigaction(SIGCHLD, &act, 0)) {
        fprintf(stderr, "Failed to install sigaction handler for SIGCHLD\n");
        exit(1);
    }

    if (setuid(geteuid())) { // become fully root
        fprintf(stderr, "Failed to setuid(%d)\n", geteuid());
        exit(1);
    }

    const int origProcessNameSize = strlen(*argv);
    renameThisProcess(*argv, origProcessNameSize, "dispatcher");

    /*
     * Loop forever and wait for messages via our pipe, from the mainListener(),
     * then we run those with the settings provided in the message.
     * This keeps the amount of logic that runs under privs to the bare minimum,
     * notice that the Runner will call exec() as the original (non-priviledged) user.
     */
    char buf[Message::MAX_SIZE + 2];
    uint32_t offset = 0;
    while (true) {
        ssize_t amount = read(inputId, buf + offset, sizeof(buf) - offset);
        if (amount == -1) {
            if (errno == EINTR) // try again.
                continue;
            fprintf(stderr, "Failed to read from pipe %d\n", errno);
            exit(0);
        }
        if (amount == 0) // pipe is closed, we should shutdown
            return;
        offset += amount;
        if (offset < 2)
            continue;
        uint32_t messageSize = static_cast<uint8_t>(buf[1]);
        messageSize = messageSize << 8;
        messageSize += static_cast<uint8_t>(buf[0]);
        if (messageSize > Message::MAX_SIZE || messageSize == 0) {
            fprintf(stderr, "Runner: invalid message size %u\n", messageSize);
            offset = 0; // discard
            continue;
        }
        if (offset < messageSize + 2)
            continue;

        try {
            Message message(buf + 2, offset - 2);
#ifdef DEBUG_MESSAGE
            message.printFields();
#endif

            auto pipeFileName = pipeFilename(ownerUid, message);
            auto pipe = findJail(occupiedJailcells, pipeFileName);
            bool sendSuccess = false;
            if (pipe) {
                size_t amount = fwrite(buf + 2, 1, offset - 2, pipe);
                if (amount == offset - 2)
                    sendSuccess = true;
                fclose(pipe);
            }

            if (!sendSuccess) {
                occupiedJailcells.push_back(pipeFileName);
                Runner runner(message, outputId);
                runner.setProcessName(*argv, origProcessNameSize);
                runner.setOwnerUserId(ownerUid);
                runner.addPipe(inputId);
                runner.run();
            }
        } catch (const std::exception &e) {
            auto len = strlen(e.what());
            write(outputId, e.what(), len + 1);
        }
        offset = 0;
    }
}

/*
 The listener listens on dbus and does the filesystem stuff and defers
 to the runner to actually start an app
*/
static void mainListener(int x, char **y, int inputId, int outputId)
{
    // drop priviledges
    const int targetUid = getuid(); // the real user, not root.
    if (setuid(targetUid) != 0) { // this is needed for DBUS to connect.
        fprintf(stderr, "Failed to change to user %d\n", targetUid);
        return;
    }
    if (seteuid(targetUid) != 0) {
        fprintf(stderr, "Failed to change to e-user %d\n", targetUid);
        return;
    }

    QApplication app(x, y);
    app.setApplicationName("isolation-runner");
    app.setQuitOnLastWindowClosed(false); // avoid magic behavior

    std::srand(std::time(0));
    QCommandLineParser parser;
    parser.addHelpOption(); // allows users to get an overview of options
    QCommandLineOption dir(QStringList() << "rulesdir" << "d", "Directory where the rules are", "PATH");
    parser.addOption(dir);
    parser.process(app);
    IsolationManager manager(inputId, outputId);
    if (parser.isSet(dir))
        manager.setRulesDir(parser.value(dir));
    app.exec();
}

int main(int x, char **y)
{
    if (geteuid() != 0) {
        fprintf(stderr, "Should be installed suid root\n");
        return 1;
    }
    if (getuid() == 0) {
        fprintf(stderr, "Do not run as root\n");
        return 1;
    }
    const char *homedir = getenv("HOME");
    if (homedir == nullptr) {
        fprintf(stderr, "Env var HOME missing\n");
        return 1;
    }
    if (chdir(homedir) != 0) {
        fprintf(stderr, "Env var HOME invalid\n");
        return 1;
    }
    if (chdir(".local/jails") != 0) {
        // if fail, create and then try to enter it again.
        if (mkdir(".local", 0700) == 0)
            chown(".local", getuid(), getgid());
        if (mkdir(".local/jails", 0700) == 0)
            chown(".local/jails", getuid(), getgid());
        if (chdir(".local/jails") != 0) {
            fprintf(stderr, "Failed to find jails dir\n");
            return 1;
        }
    }

    if (getgid() == 0 || getegid() == 0) {
        fprintf(stderr, "Warn: group is root, likely setup problem\n");
    }

    int pUp[2]; // from runner to listener
    if (pipe(pUp) == -1) {
        return 1;
    }
    int pDown[2]; // from listener to runner
    if (pipe(pDown) == -1) {
        return 1;
    }

    pid_t pid = fork();
    if (pid < 0) { // fork failed
        return 1;
    }
    if (pid) { // parent
        close(pUp[1]);
        close(pDown[0]);
        mainListener(x, y, pUp[0], pDown[1]);
        close(pUp[0]);
        close(pDown[1]);
    }
    else {
        close(pUp[0]);
        close(pDown[1]);
        mainRunner(y, pDown[0], pUp[1]);
        close(pUp[1]);
        close(pDown[0]);
    }
    return 0;
}