#define _GNU_SOURCE

#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <sodium.h>
#include <stdbool.h>
#include "common.c"

#define ENVVAR_SK_NAME "LIBEXIM_SECRETKEY"
#define ENVVAR_PK_NAME "LIBEXIM_PUBLICKEY"

void print_usage(char *progname) {
    printf("Usage: %s [OPTIONS] [CIPHERTEXT]\n\n", progname);
    printf("Secret and public key:\n");
    printf("  -s, --secret-key SECRETKEY   Secret key (base64-encoded)\n");
    printf("  -p, --public-key PUBLICKEY   Public key (base64-encoded)\n");
    printf("  -S, --secret-key-file FILE   Read secret key (raw) from file FILE (use - for stdin)\n");
    printf("  -P, --public-key-file FILE   Read public key (raw) from file FILE (use - for stdin)\n");
    printf("\n");
    printf("The environment variables %s and %s may contain base64-encoded secret/public keys. \n", ENVVAR_SK_NAME,
           ENVVAR_PK_NAME);
    printf("\n");
    printf("Select input:\n");
    printf("  -f, --infile FILE            Decrypt contents of the first line of file FILE (use - for stdin)\n");
    printf("\n");
    printf("Output:\n");
    printf("  -n, --no-newline             Do not append a newline to the output\n");
    printf("\n");
    printf("Keys in arguments and environment variables are expected to be base64 encoded (as produced by the library).\n");
    printf("Keys in files need to be raw bytes with no encoding, ciphertext should always be base64-encoded.\n");
    printf("\n");
}

typedef enum {
    NONE = 0,
    SK = 1,
    PK = 2,
    SKENV = 4,
    PKENV = 8,
    SKFILE = 16,
    PKFILE = 32,
    INFILE = 64,
    INSTRING = 128
} seen_sb_args;

int main(int argc, char *argv[]) {
    char *prog_basename = basename(argv[0]);
    int opt = 0;
    unsigned char *secretkey, *publickey;
    size_t secretkey_len = 0, publickey_len = 0;
    char *b64cipherstring;
    unsigned char *cipherstring;
    size_t cipherstring_len;
    bool add_newline = true;

    seen_sb_args publickey_mode = NONE;
    seen_sb_args secretkey_mode = NONE;
    seen_sb_args input = NONE;

    if (sodium_init() < 0) {
        fputs("Unable to initialize libsodium", stderr);
        exit(128);
    }

    // define arguments
    const char *shortargs = "s:p:S:P:f:nh";
    static struct option long_options[] = {
            {"secret-key",      required_argument, NULL, 's'},
            {"public-key",      required_argument, NULL, 'p'},
            {"secret-key-file", required_argument, NULL, 'S'},
            {"public-key-file", required_argument, NULL, 'P'},
            {"infile",          required_argument, NULL, 'f'},
            {"no-newline",      required_argument, NULL, 'n'},
            {"help",            no_argument,       NULL, 'h'},
            {0,                 0,                 0,    0}
    };

    // check environment for LIBEXIM_SECRETKEY variable
    char *sk_env = getenv(ENVVAR_SK_NAME);
    if (sk_env != NULL && strlen(sk_env) > 0) {
        if (base64_decode_string(sk_env, &secretkey, &secretkey_len) != 0) {
            fprintf(stderr, "[ERROR] Unable to base64-decode secret key.\n\n");
            exit(EXIT_FAILURE);
        }
        secretkey_mode |= SKENV;
    }

    // check environment for LIBEXIM_PUBLICKEY variable
    char *pk_env = getenv(ENVVAR_PK_NAME);
    if (pk_env != NULL && strlen(pk_env) > 0) {
        if (base64_decode_string(pk_env, &publickey, &publickey_len) != 0) {
            fprintf(stderr, "[ERROR] Unable to base64-decode public key.\n\n");
            exit(EXIT_FAILURE);
        }
        publickey_mode |= PKENV;
    }

    // parse arguments
    int long_index = 0;
    while ((opt = getopt_long(argc, argv, shortargs,
                              long_options, &long_index)) != -1) {
        switch (opt) {
            case 's':
                if (base64_decode_string(optarg, &secretkey, &secretkey_len) != 0) {
                    fprintf(stderr, "[ERROR] Unable to base64-decode secret key.\n\n");
                    exit(EXIT_FAILURE);
                }
                secretkey_mode |= SK;
                break;
            case 'p':
                if (base64_decode_string(optarg, &publickey, &publickey_len) != 0) {
                    fprintf(stderr, "[ERROR] Unable to base64-decode public key.\n\n");
                    exit(EXIT_FAILURE);
                }
                publickey_mode |= PK;
                break;
            case 'S':
                secretkey = (unsigned char *) read_password_file(optarg, crypto_box_SECRETKEYBYTES, &secretkey_len);
                secretkey_mode |= SKFILE;
                break;
            case 'P':
                publickey = (unsigned char *) read_password_file(optarg, crypto_box_PUBLICKEYBYTES, &publickey_len);
                publickey_mode |= PKFILE;
                break;
            case 'f':
                b64cipherstring = read_first_line(optarg);
                input |= INFILE;
                break;
            case 'n':
                add_newline = false;
                break;
            case 'h':
                print_usage(prog_basename);
                exit(EXIT_SUCCESS);
                break;
        }
    }

    // read first non-option argument as ciphertext if present
    if (optind < argc) {
        size_t b64cipherstring_len = strlen(argv[optind]);
        b64cipherstring = malloc(b64cipherstring_len);
        sodium_memzero(b64cipherstring, b64cipherstring_len);
        strncpy(b64cipherstring, argv[optind], b64cipherstring_len);
        input |= INSTRING;
    }

    // check if a secret key was provided
    if (secretkey_mode == NONE) {
        fprintf(stderr, "[ERROR] Please specify a secret key.\n\n");
        print_usage(prog_basename);
        exit(EXIT_FAILURE);
    }
    // check if the secret key has the correct size
    if (secretkey_len != crypto_box_SECRETKEYBYTES) {
        fprintf(stderr, "[ERROR] Secret key has wrong size %zu; expected %d.\n\n", secretkey_len,
                crypto_box_SECRETKEYBYTES);
        exit(EXIT_FAILURE);
    }

    // check if a public key was provided
    if (publickey_mode == NONE) {
        fprintf(stderr, "[ERROR] Please specify a public key.\n\n");
        print_usage(prog_basename);
        exit(EXIT_FAILURE);
    }
    // check if the public key has the correct size
    if (publickey_len != crypto_box_PUBLICKEYBYTES) {
        fprintf(stderr, "[ERROR] Secret key has wrong size %zu; expected %d.\n\n", publickey_len,
                crypto_box_PUBLICKEYBYTES);
        exit(EXIT_FAILURE);
    }

    // check if a ciphertext was provided
    if (input == NONE) {
        fprintf(stderr, "[ERROR] Please specify a ciphertext source.\n\n");
        print_usage(prog_basename);
        exit(EXIT_FAILURE);
    }
    // base64-decode ciphertext
    if (base64_decode_string(b64cipherstring, &cipherstring, &cipherstring_len) != 0) {
        fprintf(stderr, "[ERROR] Unable to base64-decode ciphertext.\n\n");
        exit(EXIT_FAILURE);
    }
    // prepare buffer for cleartext
    size_t cleartext_len = cipherstring_len - crypto_box_SEALBYTES;
    unsigned char *cleartext = (unsigned char *) malloc(cleartext_len + 1);
    sodium_memzero(cleartext, cleartext_len + 1);

    // decrypt message
    if (crypto_box_seal_open(cleartext, cipherstring, cipherstring_len, publickey, secretkey) != 0) {
        fprintf(stderr, "[ERROR] Unable to decrypt ciphertext.\n\n");
        exit(EXIT_FAILURE);
    }

    // print cleartext to stdout
    if (add_newline == true) {
        fprintf(stdout, "%s\n", (const char *) cleartext);
    } else {
        fprintf(stdout, "%s", (const char *) cleartext);
    }
}