/*
 * The functions in this file are responsible for most of the special
 * functionality of the purr/meow utilities, which is encrypted pastes
 * that the server can't access. They can be configured with the RANDOMIZE_IV,
 * ENCODE_BASE_64 and DECODE_BASE_64 macros from purr.h.
 *
 * RANDOMIZE_IV: IV used for AES encryption is randomized or all zero.
 *
 * ENCODE_BASE_64: the binary data generated by the encryption process is
 * encoded as base64.
 *
 * DECODE_BASE_64: the data returned by the server is decoded, assuming its
 * encoding is base64.
 *
 * Having all these macros defined is the default mode of operation, and what's
 * expected by the other PurritoBin clients.
 */

#define _DEFAULT_SOURCE /* getentropy */
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <assert.h>
#include <sys/random.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <bearssl.h>

#include "purr.h"
#include "mmap_file.h"

#define PEM_HEADER "TOP-SECRET"
// 5 slashes + "BEGIN " + <header> + 5 slashes + LR
#define PEM_HEADER_BEGIN_SIZE (sizeof PEM_HEADER - 1 + 5 + sizeof("BEGIN ") - 1 + 5 + 1)
// LR + 5 slashes + "END " + <header> + 5 slashes
#define PEM_HEADER_END_SIZE (1 + sizeof PEM_HEADER - 1 + 5 + sizeof("END ") - 1 + 5)

/*
 * This function takes an mmap_file struct, and returns an encrypted mmap_file from it.
 * Args:
 *   file: mmap_file for the input file
 *   keyp: will receive the newly generated random key
 *   ivp: will receive the newly generated random IV (if enabled in purr.h)
 */
struct mmap_file encrypt_mmap(struct mmap_file file, uint8_t **keyp, uint8_t **ivp)
{
    off_t file_size = file.size;
    ssize_t blocks = file_size / br_aes_big_BLOCK_SIZE;
    if (blocks * br_aes_big_BLOCK_SIZE < file_size) blocks++;
    file_size = blocks * br_aes_big_BLOCK_SIZE;
    ssize_t padding = file_size - file.size;

    struct mmap_file rv = {.size = file_size, .prot = PROT_MEM, .flags = MAP_MEM};

    uint8_t *key = calloc(KEY_LEN, 1);
    uint8_t *iv = calloc(IV_LEN, 1);
    uint8_t iv_throwaway[IV_LEN] = { 0 };
    if (key == NULL || iv == NULL) {
        perror("allocation failure");
        return rv;
    }

    if (getentropy(key, KEY_LEN) < 0) {
        perror("getentropy()");
        return rv;
    }
    #ifdef RANDOMIZE_IV
    if (getentropy(iv, IV_LEN) < 0) {
        perror("getentropy()");
        return rv;
    }
    memcpy(iv_throwaway, iv, IV_LEN);
    #endif /* RANDOMIZE_IV */

    if (!allocate_mmap(&rv)) {
        return rv;
    }

    memcpy(rv.data, file.data, file.size);
    // PKCS7 padding
    memset(rv.data + file.size, padding, padding);

    br_aes_big_cbcenc_keys br = { 0 };
    br_aes_big_cbcenc_init(&br, key, KEY_LEN);
    br_aes_big_cbcenc_run(&br, iv_throwaway, rv.data, file_size);

    #ifdef ENCODE_BASE_64
    struct mmap_file rv_64 = {.prot = PROT_MEM, .flags = MAP_MEM};

    // default mode -> PEM object with 76 characters per line
    size_t encsize = br_pem_encode(NULL, rv.data, rv.size, PEM_HEADER, 0);
    char *pem_enc = malloc(encsize + 1);
    if (pem_enc == NULL) {
        return rv_64;
    }
    br_pem_encode(pem_enc, rv.data, rv.size, PEM_HEADER, 0);

    int newlines = 0;
    char *new_line_search = pem_enc;
    while((new_line_search = strchr(new_line_search, '\n'))) {
        new_line_search++;
        newlines++;
    }
    if (newlines < 2) {
        fputs("encoding error!", stderr);
        return rv_64;
    } else {
        newlines -= 2;
    }

    rv_64.size = encsize - newlines - PEM_HEADER_BEGIN_SIZE - PEM_HEADER_END_SIZE;
    if (!allocate_mmap(&rv_64)) {
        return rv_64;
    }
    for (int i = 0; i < newlines; i++) {
        memcpy(rv_64.data + i * 76,
               pem_enc + PEM_HEADER_BEGIN_SIZE + i * (76 + 1),
               (i == newlines) ? rv_64.size % 76 : 76);
    }

    free(pem_enc);
    free_mmap(&rv);
    rv = rv_64;
    #endif /* ENCODE_BASE_64 */

    free_mmap(&file);

    // pass pointers to caller
    *keyp  = key;
    *ivp = iv;

    return rv;
}

#ifdef DECODE_BASE_64
struct b64_decoding_struct {
    uint8_t *data;
    size_t n, size;
    bool dead;
};

enum b64_decoding_status {
B64_BEGIN,
B64_MAIN,
B64_END
};

static void push_decoded_b64(void *dest_ctx, const void *src, size_t len)
{
    struct b64_decoding_struct *ct = dest_ctx;
    if (ct->dead) return;

    if (ct->n + len > ct->size) {
        if (ct->size == 0) ct->size = 64;

        while (ct->n + len > (ct->size *= 2));

        ct->data = realloc(ct->data, ct->size);
        if (ct->data == NULL) {
            perror("realloc()");
            ct->dead = true;
            return;
        }
    }

    memcpy(ct->data + ct->n, src, len);
    ct->n += len;
}
#endif /* DECODE_BASE_64 */

/*
 * This function takes an mmap_file struct, and returns a decrypted buffer from it.
 * Args:
 *   file: mmap_file for the input file
 *   keyp: the key used to encrypt the paste
 *   ivp: the iv used to encrypt the paste
 */
struct mmap_file decrypt_mmap(struct mmap_file file, const uint8_t *key, const uint8_t *iv)
{
    struct mmap_file rv = {.size = file.offset, .prot = PROT_MEM, .flags = MAP_MEM};

    #ifdef DECODE_BASE_64
    br_pem_decoder_context pem;
    br_pem_decoder_init(&pem);
    struct b64_decoding_struct dec = { 0 };
    br_pem_decoder_setdest(&pem, push_decoded_b64, &dec);

    const char *b64_begin_header = "-----BEGIN TOP-SECRET-----\n";
    const size_t beginlen_c = strlen(b64_begin_header);
    size_t beginlen = beginlen_c;

    static const char *b64_end_header = "\n-----END TOP-SECRET-----";
    const size_t endlen_c = strlen(b64_end_header);
    size_t endlen = endlen_c;

    size_t wlen = file.offset;

    enum b64_decoding_status dec_step = B64_BEGIN;
    while (1) {
        bool should_break = false;

        switch (dec_step) {
            size_t pushed;

            case B64_BEGIN:
                pushed = br_pem_decoder_push(
                    &pem,
                    b64_begin_header + (beginlen_c - beginlen),
                    beginlen);
                beginlen -= pushed;
                if (beginlen == 0) dec_step = B64_MAIN;
                break;
            case B64_MAIN:
                pushed = br_pem_decoder_push(
                    &pem,
                    file.data + (file.offset - wlen),
                    wlen);
                wlen -= pushed;
                if (wlen == 0) dec_step = B64_END;
                break;
            case B64_END:
                pushed = br_pem_decoder_push(
                    &pem,
                    b64_end_header + (endlen_c - endlen),
                    endlen);
                endlen -= pushed;
                if (endlen == 0) should_break = true;
                break;
        }

        switch (br_pem_decoder_event(&pem)) {
            case BR_PEM_ERROR:
                fputs("base64 decoding error!\n", stderr);
                dec.dead = true;
                // fall through
            case BR_PEM_END_OBJ:
                should_break = true;
                // fall through
            case BR_PEM_BEGIN_OBJ:
            case 0:
                break;
        }
        if (should_break) break;
    }
    if (dec.dead) {
        return rv;
    }
    rv.size = dec.n;
    #endif /* DECODE_BASE_64 */

    if (rv.size % br_aes_big_BLOCK_SIZE) {
        fputs("bad input size!\n", stderr);
        return rv;
    }
    if (!allocate_mmap(&rv)) {
        return rv;
    }

    #ifdef DECODE_BASE_64
    write_into_mmap(&rv, dec.data, rv.size);
    free(dec.data);
    #else
    write_into_mmap(&rv, file.data, rv.size);
    #endif /* DECODE_BASE_64 */

    free_mmap(&file);

    uint8_t iv_throwaway[IV_LEN];
    memcpy(iv_throwaway, iv, IV_LEN);

    br_aes_big_cbcdec_keys br = { 0 };
    br_aes_big_cbcdec_init(&br, key, KEY_LEN);
    br_aes_big_cbcdec_run(&br, iv_throwaway, rv.data, rv.offset);

    // remove padding - only knows PKCS7
    int padding = rv.data[rv.size - 1];
    if (padding < br_aes_big_BLOCK_SIZE) {
        // data might contain padding
        bool found_padding = true;
        for (int i = 0; i < padding; i++) {
            if (rv.data[rv.offset - 1 - i] != padding) {
                found_padding = false;
            }
        }
        if (found_padding) {
            rv.offset -= padding;
        }
    }

    return rv;
}