#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 get_self_filepath(char *buffer, size_t size) {
    ssize_t bytes_written = readlink("/proc/self/exe", buffer, size - 1);
    if(bytes_written == -1)
        return 0;

    buffer[bytes_written] = '\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;

    snprintf(path, sizeof(path), "%s/.local", home);
    err = mkdir(path, S_IRWXU);
    if(err == -1 && errno != EEXIST)
        return 0;

    snprintf(path, sizeof(path), "%s/.local/bin", home);
    err = mkdir(path, S_IRWXU);
    if(err == -1 && errno != EEXIST)
        return 0;

    return 1;
}

static int create_local_kms_server_proxy_file(const char *source_path, const char *dest_path) {
    int in_fd = -1;
    int out_fd = -1;
    int res = 0;

    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(dest_path, 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(dest_path, 1);

    return res;
}

int main(int argc, char **argv) {
    if(argc != 4)
        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(!get_self_filepath(self_path, sizeof(self_path)))
        return 1;

    char kms_server_proxy_local_filepath[PATH_MAX];
    snprintf(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), "%s/.local/bin/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)) {
        if(strcmp(self_path, kms_server_proxy_local_filepath) != 0) {
            const char *args[] = { kms_server_proxy_local_filepath, 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(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);
    }
}