#include "common.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <sys/wait.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/user.h>
#include <sys/syscall.h>
#include <linux/limits.h>



char *get_pid_rel_path(execontrol_pid_t pid, const char *path, size_t path_len) {
    char *pid_rel_path = malloc(sizeof("/proc/99999999/cwd/")+path_len);
    size_t pid_rel_path_len = sprintf(pid_rel_path, "/proc/%d/cwd/", pid);
    memcpy(pid_rel_path+pid_rel_path_len, path, path_len);
    return pid_rel_path;
}

int is_rel_path(const char *path) {
    return path[0] == '.';
}

struct execontrol_response can_file(execontrol_cmd_t cmd, execontrol_pid_t pid, const char *filename, size_t filename_len, int *control_pipe) {
    struct execontrol_request request = {
        .pid = pid,
        .cmd = cmd
    };
    do_checked(write, control_pipe[1], &request, sizeof(request));
    do_checked(write, control_pipe[1], filename, filename_len+1);
    struct execontrol_response response;
    do_checked(read, control_pipe[0], &response, sizeof(response));
    return response;
}

struct execontrol_response can_file_maybe_relative(execontrol_cmd_t cmd, execontrol_pid_t pid, const char *filename, size_t filename_len, int *control_pipe) {
    struct execontrol_response response;
    if (is_rel_path(filename)) {
        char *absolute_path = get_pid_rel_path(pid, filename, filename_len);
        response = can_file(cmd, pid, absolute_path, strlen(absolute_path), control_pipe);
        free(absolute_path);
    } else {
        response = can_file(cmd, pid, filename, filename_len, control_pipe);
    }
    return response;
}


size_t read_string(char *buffer, size_t buffer_len, execontrol_pid_t pid, unsigned long long addr) {
    // Read string
    size_t string_len;
    for (string_len = 0; string_len != buffer_len-1; string_len++) {
        buffer[string_len] = ptrace(PTRACE_PEEKDATA, pid, addr+string_len);
        if (buffer[string_len] == '\0') {
            break;
        }
        if (buffer[string_len] < 0) {
            buffer[string_len] = '\0';
            break;
        }
    }
    buffer[buffer_len-1] = '\0';
    return string_len;
}

void run(execontrol_pid_t pid, int *sync, int *control_pipe) {
    // Start ptrace
    do_checked(ptrace, PTRACE_SEIZE, pid, 0, 0);
    do_checked(ptrace, PTRACE_INTERRUPT, pid, 0, 0);
    do_checked(waitpid, pid, 0, 0);
    do_checked(ptrace, PTRACE_SETOPTIONS, pid, 0, PTRACE_O_EXITKILL);
    // Run
    int lastsig = 0;
    char syscall_result = 0;
    for (;;) {
        int status;
        // Wait for next event
        do_checked(ptrace, PTRACE_SYSCALL, pid, 0, lastsig);
        do_checked(waitpid, pid, &status, 0);
        // Handle exit
        if (WIFEXITED(status)) {
            exit(WEXITSTATUS(status));
        }
        // Handle signals
        if (WIFSTOPPED(status)) {
            lastsig = WSTOPSIG(status);
            // Handle trap
            if (lastsig == SIGTRAP) {
                lastsig = 0;
                // Must be a syscall; handle it
                struct user_regs_struct regs;
                do_checked(ptrace, PTRACE_GETREGS, pid, 0, &regs);
                switch (regs.orig_rax) {
                case SYS_execve: {
                    if (!syscall_result) {
                        // Read filename
                        char filename[PATH_MAX];
                        size_t filename_len = read_string(filename, sizeof(filename), pid, regs.rdi);
                        // Check if allowed to execute
                        execontrol_errno_t potential_errno = can_file_maybe_relative(EXECONTROL_CAN_EXEC, pid, filename, filename_len, control_pipe).error;
                        // Don't pass through if not allowed to
                        if (potential_errno <= 0) {
                            regs.orig_rax = -1;
                            regs.rax = potential_errno?potential_errno:-ECANCELED;
                        }
                    }
                } break;
                case SYS_clone: {
                    if (syscall_result) {
                        // Attach to it
                        execontrol_pid_t npid = regs.rax;
                        do_checked(write, sync[1], &npid, sizeof(npid));
                    }
                } break;
                case SYS_execveat: {
                    regs.orig_rax = -1;
                    regs.rax = -ENOTSUP;
                } break;
                }
                // Toggle syscall_result flag
                syscall_result = !syscall_result;
                do_checked(ptrace, PTRACE_SETREGS, pid, 0, &regs);
            }
        }
    }
}

int main(int argc, char **argv) {
    // Check arguments
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <pipe> <executable> [args...]\n", argv[0]);
        return -EINVAL;
    }
    // Get arguments
    char *control_pipe_path = argv[1],
         *executable = argv[2],
        **executable_argv = argv+2;
    size_t control_pipe_path_len = strlen(control_pipe_path);
    // Create control pipe
    char *control_pipe_full_path[] = {
        malloc(control_pipe_path_len+sizeof(".in")),
        malloc(control_pipe_path_len+sizeof(".out"))
    };
    memcpy(control_pipe_full_path[0], control_pipe_path, control_pipe_path_len);
    memcpy(control_pipe_full_path[1], control_pipe_path, control_pipe_path_len);
    memcpy(control_pipe_full_path[0]+control_pipe_path_len, ".in", sizeof(".in"));
    memcpy(control_pipe_full_path[1]+control_pipe_path_len, ".out", sizeof(".out"));
    mkfifo(control_pipe_full_path[0], 0600);
    mkfifo(control_pipe_full_path[1], 0600);
    int control_pipe_write_end_ = do_checked(open, control_pipe_full_path[1], O_WRONLY); // We need to do open the files in opposite order to not get stuck
    int control_pipe[] = {
        do_checked(open, control_pipe_full_path[0], O_RDONLY),
        control_pipe_write_end_
    };
#   ifdef NDEBUG
    unlink(control_pipe_full_path[0]);
    unlink(control_pipe_full_path[1]);
#   endif
    free(control_pipe_full_path[0]);
    free(control_pipe_full_path[1]);
    // Create synchronization buffer
    int sync[2];
    do_checked(pipe, sync);
    char syncbuf[1];
    // Start process
    execontrol_pid_t pid = fork();
    if (pid == 0) {
        do_checked(write, sync[1], syncbuf, sizeof(syncbuf));
        // Close pipes
        close(sync[0]);
        close(sync[1]);
        close(control_pipe[0]);
        close(control_pipe[1]);
        // Run given program
        execvp(executable, executable_argv);
        perror(executable);
        return -errno;
    }
    do_checked(read, sync[0], syncbuf, sizeof(syncbuf));
    // Run
    if (fork() == 0) {
        run(pid, sync, control_pipe);
        exit(0);
    }
    // Attach to all further processes
    if (fork() == 0) {
        for(;;) {
            execontrol_pid_t npid;
            do_checked(read, sync[0], &npid, sizeof(npid));
            if (fork() == 0) {
                run(npid, sync, control_pipe);
                return 0;
            }
        }
    }
    // Wait for process to exit
    waitpid(pid, NULL, 0);
}