#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <sys/random.h>

#include <monocypher.h>
#include "epic_drm.h"

#define log(what) printf("[epic_drm] " what "\n")

/**
 * sends a protected command to the client
 */
void get_cmd(struct command_channel* chan, const uint8_t* key) {
    log("sending encrypted DRM command");
    uint8_t cmd[CMD_SIZE] = "echo 'this shell command has been protected with Epic DRM(tm)'";
    getrandom(chan->data.nonce, 24, 0);
    crypto_lock(chan->data.mac, chan->data.exec_cmd, key, chan->data.nonce, cmd, CMD_SIZE);
    chan->command = SUCCESS;
}

/**
 * unprotects and executes data from the client
 */
void execute_cmd(struct command_channel* chan, const uint8_t* key) {
    log("decrypting command");
    uint8_t cmd[CMD_SIZE];
    int res = crypto_unlock(cmd, key, chan->data.nonce, chan->data.mac, chan->data.exec_cmd, CMD_SIZE);
    if (res != 0) {
        log("decryption error");
        chan->command = ERROR;
    } else {
        log("executing command");
        system((const char*)cmd);
        chan->command = SUCCESS;
    }
}

int main() {
    // create random key
    uint8_t key[16];
    log("generating random key");
    getrandom(key, sizeof(key), 0);

    // initialize memory
    log("opening shm");
    shm_unlink(SHM_NAME);
    int shm = shm_open(SHM_NAME, O_RDWR | O_CREAT | O_EXCL, 0644);
    ftruncate(shm, SHM_SIZE);
    struct command_channel* chan = mmap(NULL, SHM_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, shm, 0);

    sem_init(&chan->sem_c2s, 1, 0);
    sem_init(&chan->sem_s2c, 1, 0);

    // process commands
    log("ready");
    while (true) {
        sem_wait(&chan->sem_c2s);

        switch (chan->command) {
            case GET_CMD:
                get_cmd(chan, key);
                break;
            case EXECUTE_CMD:
                execute_cmd(chan, key);
                break;
            default:
                log("unexpected command");
        }

        sem_post(&chan->sem_s2c);
    }
}