#define _GNU_SOURCE
#include <sys/wait.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <sched.h>
#include <sys/syscall.h>
#include <linux/sched.h>
#include <stdarg.h>
#include <time.h>
#include <sys/time.h>
#include <limits.h>
#include <signal.h>
#include <stdbool.h>

#include "util.h"

struct params {
    char* shared_folder;
    int fd[2];
    int out_fd;
    char **argv;
};

struct limits {
    size_t memory;
    int core;
    int time;
    int processes;
};

struct killparams{
    int time;
    int pid;
};

#include "ns_exec.c"
#include "cgroup_prepare.c"

uint64_t time_start;

// get time in milliseconds
uint64_t getmstime() {
    struct timeval tm;
    gettimeofday(&tm, NULL);
    return (uint64_t) tm.tv_sec * 1000 + (uint64_t) tm.tv_usec / 1000;
}

// time limit realisation
static int killafter(void *arg) {
    // die when parent dies
    if (prctl(PR_SET_PDEATHSIG, SIGKILL)) die("cannot PR_SET_PDEATHSIG for child process: %m\n");
    struct killparams* params = (struct killparams*) arg;
    struct timespec tw = {params -> time / 1000, ((params->time) % 1000+5) * 1000000};
    struct timespec tr;
    nanosleep(&tw, &tr);
    kill(params->pid, SIGKILL);
    return 0;
}

#define STACK_SIZE 1000000
static char nmstack[STACK_SIZE];
static char killstack[STACK_SIZE];

static void parse_args(int argc, char **argv, struct params *params, struct limits *limits){
    if (argc < 7) {
        puts("usage:\n starter <max processes number> <core_id> <memory_amount> <time limit in ms> <shared folder> <command, arg1,arg2,...>");
        exit(0);
    }
    argc--; argv++;
    limits->processes = atoi(argv[0]);
    argc--; argv++;
    limits->core = atoi(argv[0]);
    argc--; argv++;
    limits->memory = atoi(argv[0]);
    argc--; argv++;
    limits->time = atoi(argv[0]);
    argc--; argv++;
    params->shared_folder = argv[0];
    argc--; argv++;

    params->argv = argv;
}

// setup user namespace
static void prepare_userns(int pid) {
    char path[100];
    char line[100];

    int uid = 0;// root
    int gid = 0;// root
    int unprivileged_uid = 66534;// nobody
    int unprivileged_gid = 65534;// nogroup

    sprintf(path, "/proc/%d/uid_map", pid);
    sprintf(line, "0 %d 1\n1 %d 1000\n", uid, unprivileged_uid);// map root to uid 0,nobody to unprivileged uid
    write_file(path, line);

    sprintf(path, "/proc/%d/setgroups", pid);
    sprintf(line, "deny");
    write_file(path, line);

    sprintf(path, "/proc/%d/gid_map", pid);
    sprintf(line, "0 %d 1\n1 %d 1000\n", gid, unprivileged_gid);// map root to gid 0,nogroup to unprivileged gid
    write_file(path, line);
}

static void get_real_path(char* path) {
    if (realpath (argv[0], path) == 0) die("unable to resolve real path: %m");// get absolute path to executable
    for(int i = strlen(path); i > 0 && path[i] != '/';i--) path[i]=0;// cut filename to get directory name
}

int main(int argc,char** argv) {
    if(setuid(0)) die("must be run as root");
    if(setgid(0)) die("must be run as root");
    // get binary path
    char real_path[PATH_MAX];
    get_real_path(real_path);
    if(chdir(real_path)) die("unable to chdir to binary path: %m");
    // set random seed
    srand(time(NULL));
    // setup parameters
    struct params params;
    memset(&params, 0, sizeof(struct params));
    struct limits limits;
    memset(&limits, 0, sizeof(struct limits));
    parse_args(argc, argv, &params, &limits);
    prepare_cgroup(&limits);
    if (pipe(params.fd) < 0) die("can't open pipe: %m");// a pipe to report readiness
    int clone_flags = SIGCHLD | CLONE_NEWUTS | CLONE_NEWUSER | CLONE_NEWNS | CLONE_NEWPID | CLONE_NEWNET | CLONE_NEWIPC | CLONE_NEWCGROUP;
    int nsrun_pid = clone(nsrun, nmstack + STACK_SIZE, clone_flags, &params);// make new namespace
    prepare_userns(nsrun_pid);
    if (nsrun_pid < 0) die("faled to clone");
    add_to_cgroup(nsrun_pid);
    if (write(params.fd[1], "OK", 2) != 2) die("Failed to write to pipe: %m");// report readiness
    struct killparams killparams;
    killparams.time = limits.time;
    killparams.pid = nsrun_pid;
    int kill_pid = clone(killafter, killstack + STACK_SIZE, SIGCHLD, &killparams);
    time_start = getmstime();
    if (waitpid(nsrun_pid, NULL, 0) == -1) die("Failed to wait pid %d: %m\n", nsrun_pid);
    kill(kill_pid, SIGKILL);// kill killer
    int time = getmstime() - time_start;
    if(chdir(params.shared_folder)) die("Failed to chdir to shared folder:%m");
    if(time >= limits.time) {
        write_file("time", "-1");
    }
    else {
        char* timestr = NULL;
        asprintf(&timestr, "%d", time);
        write_file("time", timestr);
        free(timestr);
    }
    int memory = check_mem();
    char* memstr = NULL;
    asprintf(&memstr, "%d", memory);
    write_file("memory", memstr);
    free(memstr);
    remove_cgroup();
    return 0;
}