#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <limits.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/sendfile.h>
#include <sys/capability.h>

static int readlink_realpath(const char *filepath, char *buffer) {
    char symlinked_path[PATH_MAX];
    ssize_t bytes_written = readlink(filepath, symlinked_path, sizeof(symlinked_path) - 1);
    if(bytes_written == -1 && errno == EINVAL) {
        /* Not a symlink */
        strncpy(symlinked_path, filepath, sizeof(symlinked_path));
    } else if(bytes_written == -1) {
        return 0;
    } else {
        symlinked_path[bytes_written] = '\0';
    }

    if(!realpath(symlinked_path, buffer))
        return 0;

    return 1;
}

static int file_has_sys_admin_capability(const char *filepath) {
    cap_t cap = cap_get_file(filepath);
    if(!cap)
        return 0;

    cap_flag_value_t res = CAP_CLEAR;
    cap_get_flag(cap, CAP_SYS_ADMIN, CAP_PERMITTED, &res);
    int cap_set = res == CAP_SET;
    cap_free(cap);
    return cap_set;
}

static int file_set_sys_admin_capability(const char *filepath, int setfcap) {
    cap_t cap = cap_get_file(filepath);
    if(!cap && errno != ENODATA)
        return 0;

    if(!cap) {
        cap = cap_init();
        if(!cap)
            return 0;
    }

    const cap_value_t cap_to_set[] = { CAP_SYS_ADMIN, CAP_SETFCAP };
    int num_caps = 1;
    if(setfcap)
        num_caps = 2;

    int res = 0;
    res |= cap_set_flag(cap, CAP_EFFECTIVE, num_caps, cap_to_set, CAP_SET);
    res |= cap_set_flag(cap, CAP_PERMITTED, num_caps, cap_to_set, CAP_SET);
    /*cap_set_flag(cap, CAP_INHERITABLE, num_caps, cap_to_set, CAP_SET);*/
    res |= cap_set_file(filepath, cap);
    cap_free(cap);
    return res == 0;
}

static int create_local_kms_server_proxy_directory(const char *home) {
    char path[PATH_MAX];
    int err;

    const char *paths[] = { ".local", ".local/share", ".local/share/gpu-screen-recorder", NULL };
    for(size_t i = 0; paths[i]; ++i) {
        const char *path_part = paths[i];
        snprintf(path, sizeof(path), "%s/%s", home, path_part);
        err = mkdir(path, S_IRWXU);
        if(err == -1 && errno != EEXIST)
            return 0;
    }

    return 1;
}

static int create_local_kms_server_proxy_file_atomic(const char *source_path, const char *dest_path) {
    int in_fd = -1;
    int out_fd = -1;
    int res = 0;
    char tmp_filepath[PATH_MAX];
    snprintf(tmp_filepath, sizeof(tmp_filepath), "%s.tmp", dest_path);

    in_fd = open(source_path, O_RDONLY);
    if(in_fd == -1)
        goto done;

    struct stat st;
    if(fstat(in_fd, &st) == -1)
        goto done;

    out_fd = open(tmp_filepath, O_RDWR | O_CREAT | O_TRUNC, 0755);
    if(out_fd == -1)
        goto done;

    if(sendfile(out_fd, in_fd, NULL, st.st_size) != st.st_size)
        goto done;

    res = 1;

    done:
    if(in_fd)
        close(in_fd);
    if(out_fd)
        close(out_fd);

    if(res)
        res = file_set_sys_admin_capability(tmp_filepath, 1);

    if(res)
        rename(tmp_filepath, dest_path);

    return res;
}

int main(int argc, char **argv) {
    if(argc != 4) {
        fprintf(stderr, "usage: kms-server-proxy <initial_socket_path> <card_path> <user_homepath>\n");
        return 1;
    }

    const char *gsr_kms_server_filepath = "/var/lib/flatpak/app/com.dec05eba.gpu_screen_recorder/current/active/files/bin/gsr-kms-server";
    const char *initial_socket_path = argv[1];
    const char *card_path = argv[2];
    const char *user_homepath = argv[3];

    char self_path[PATH_MAX];
    if(!readlink_realpath("/proc/self/exe", self_path)) {
        fprintf(stderr, "failed to resolve /proc/self/exe\n");
        return 1;
    }

    char kms_server_proxy_local_filepath[PATH_MAX];
    /* Update kms-server-proxy-N to kms-server-proxy-N+1 on update (update that needs to run before this program launches itself) */
    snprintf(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), "%s/.local/share/gpu-screen-recorder/kms-server-proxy-1", user_homepath);

    if(file_has_sys_admin_capability(gsr_kms_server_filepath)) {
        const char *args[] = { gsr_kms_server_filepath, initial_socket_path, card_path, NULL };
        return execv(args[0], (char *const*)args);
    } else if(file_has_sys_admin_capability(kms_server_proxy_local_filepath)) {
        /* Need to resolve kms_server_proxy_local_filepath because /home can be a symlink to another location */
        char kms_server_proxy_local_filepath_full[PATH_MAX];
        if(!readlink_realpath(kms_server_proxy_local_filepath, kms_server_proxy_local_filepath_full)) {
            fprintf(stderr, "failed to resolve %s\n", kms_server_proxy_local_filepath);
            return 1;
        }

        /*
            Run cached ~/.local/share/gpu-screen-recorder/kms-server-proxy which has sys admin capability.
            The one in flatpak location gets capabilities overwritten on flatpak update.
        */
        if(strcmp(self_path, kms_server_proxy_local_filepath_full) != 0) {
            const char *args[] = { kms_server_proxy_local_filepath_full, initial_socket_path, card_path, user_homepath, NULL };
            return execv(args[0], (char *const*)args);
        }

        /* TODO: Remove the need for this. Instead inherit capabilities */
        if(file_set_sys_admin_capability(gsr_kms_server_filepath, 0)) {
            const char *args[] = { gsr_kms_server_filepath, initial_socket_path, card_path, NULL };
            return execv(args[0], (char *const*)args);
        }

        const char *args[] = { "pkexec", gsr_kms_server_filepath, initial_socket_path, card_path, NULL };
        return execvp(args[0], (char *const*)args);
    } else if(geteuid() == 0) { /* is current user root */
        file_set_sys_admin_capability(gsr_kms_server_filepath, 0);
        create_local_kms_server_proxy_file_atomic(self_path, kms_server_proxy_local_filepath);

        const char *args[] = { gsr_kms_server_filepath, initial_socket_path, card_path, NULL };
        return execv(args[0], (char *const*)args);
    } else {
        create_local_kms_server_proxy_directory(user_homepath);

        const char *args[] = { "pkexec", self_path, initial_socket_path, card_path, user_homepath, NULL };
        return execvp(args[0], (char *const*)args);
    }
}