#include "tweetpipe_p.h"
#include "tweetpwhash.h"
#include "util.h"

static struct option getopt_long_options[] = {
    { "help", 0, NULL, 'h' },     { "decrypt", 0, NULL, 'd' }, { "encrypt", 0, NULL, 'e' },
    { "in", 1, NULL, 'i' },       { "out", 1, NULL, 'o' },     { "pass", 1, NULL, 'p' },
    { "passfile", 1, NULL, 'P' }, { "passgen", 0, NULL, 'G' }, { NULL, 0, NULL, 0 }
};
static const char *getopt_options = "hdeGi:o:p:P:";

static void
usage(FILE *stream)
{
    fputs(
        "Usage:\n"
        "    tweetpipe -G\n"
        "    tweetpipe {-e | -d} {-p <string> | -P <file>} [-i <file>] [-o <file>]\n"
        "\n"
        "Options:\n"
        "    -G, --passgen          generate a random password\n"
        "    -e, --encrypt          encryption mode\n"
        "    -d, --decrypt          decryption mode\n"
        "    -p, --pass <password>  use <password>\n"
        "    -P, --passfile <file>  read password from <file>\n"
        "    -i, --in <file>        read input from <file>\n"
        "    -o, --out <file>       write output to <file>\n"
        "    -h, --help             print this message\n",
        stream);
    exit(2);
}

static int
file_open(const char *file, int create)
{
    int fd;

    if (file == NULL || (file[0] == '-' && file[1] == 0)) {
        return create ? STDOUT_FILENO : STDIN_FILENO;
    }
    fd = create ? open(file, O_CREAT | O_WRONLY | O_TRUNC, 0644) : open(file, O_RDONLY);
    if (fd == -1) {
        die(1, "Unable to access [%s]", file);
    }
    return fd;
}

static void
derive_key(Context *ctx, char *password, size_t password_len)
{
    static uint8_t master_key[crypto_pwhash_SALTBYTES] = { 0 };

    if (ctx->has_key) {
        die(0, "A single key is enough");
    }
    if (crypto_pwhash(ctx->key, sizeof ctx->key, (unsigned char *) password, password_len, master_key, PWHASH_OPSLIMIT) != 0) {
        die(0, "Password hashing failed");
    }
    memzero(password, password_len);
    ctx->has_key = 1;
}

static int
stream_encrypt(Context *ctx)
{
    unsigned char *const chunk_size_p = ctx->buf;
    unsigned char *const chunk_nonce  = chunk_size_p + 4;
    unsigned char *const chunk_base   = chunk_nonce + crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES;
    unsigned char *const chunk_msg    = chunk_base + crypto_secretbox_ZEROBYTES;
    unsigned char nonce[crypto_secretbox_NONCEBYTES];
    uint64_t             chunk_id;
    ssize_t              max_chunk_size;
    ssize_t              chunk_size;

    assert(ctx->sizeof_buf >= 4 + crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES);
    max_chunk_size = ctx->sizeof_buf - 4 - crypto_secretbox_NONCEBYTES + crypto_secretbox_BOXZEROBYTES - crypto_secretbox_ZEROBYTES;
    assert(max_chunk_size <= 0x7fffffff);
    chunk_id = 0;
    while ((chunk_size = safe_read_partial(ctx->fd_in, chunk_msg, max_chunk_size)) >= 0) {
        STORE32_LE(chunk_size_p, (uint32_t) chunk_size);
        memzero(chunk_nonce, crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES);
        randombytes(nonce, crypto_secretbox_NONCEBYTES);
        if (crypto_secretbox(chunk_base, chunk_base, chunk_size + crypto_secretbox_ZEROBYTES, nonce, ctx->key) !=
            0) {
            die(0, "Encryption error");
        }
        memcpy(chunk_nonce, nonce, crypto_secretbox_NONCEBYTES);
        if (safe_write(ctx->fd_out, chunk_size_p, 4 + crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES + chunk_size,
                       -1) < 0) {
            die(1, "write()");
        }
        if (chunk_size == 0) {
            break;
        }
        chunk_id++;
    }
    if (chunk_size < 0) {
        die(1, "read()");
    }
    return 0;
}

static int
stream_decrypt(Context *ctx)
{
    unsigned char *const chunk_size_p = ctx->buf;
    unsigned char *const chunk_nonce  = chunk_size_p + 4;
    unsigned char *const chunk_base   = chunk_nonce + crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES;
    unsigned char *const chunk_msg    = chunk_base + crypto_secretbox_ZEROBYTES;
    unsigned char nonce[crypto_secretbox_NONCEBYTES];
    uint64_t             chunk_id;
    ssize_t              readnb;
    ssize_t              max_chunk_size;
    ssize_t              chunk_size;

    assert(ctx->sizeof_buf >= 4 + crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES);
    max_chunk_size = ctx->sizeof_buf - 4 - crypto_secretbox_NONCEBYTES + crypto_secretbox_BOXZEROBYTES - crypto_secretbox_ZEROBYTES;
    assert(max_chunk_size <= 0x7fffffff);
    chunk_id = 0;
    while ((readnb = safe_read(ctx->fd_in, chunk_size_p, 4)) == 4) {
        chunk_size = LOAD32_LE(chunk_size_p);
        if (chunk_size > max_chunk_size) {
            die(0, "Chunk size too large ([%zd] > [%zd])", chunk_size, max_chunk_size);
        }
        if (safe_read(ctx->fd_in, chunk_nonce, crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES + chunk_size) !=
            crypto_secretbox_NONCEBYTES - crypto_secretbox_BOXZEROBYTES + crypto_secretbox_ZEROBYTES + chunk_size) {
            die(0, "Chunk too short ([%zd] bytes expected)", chunk_size);
        }
        memcpy(nonce, chunk_nonce, crypto_secretbox_NONCEBYTES);
        memzero(chunk_nonce, crypto_secretbox_NONCEBYTES);
        if (crypto_secretbox_open(chunk_base, chunk_base, chunk_size + crypto_secretbox_ZEROBYTES, nonce, ctx->key) != 0) {
            printf("Unable to decrypt chunk #%" PRIu64 " - ", chunk_id);
            if (chunk_id == 0) {
                die(0, "Wrong password or key?");
            } else {
                die(0, "Corrupted or incomplete file?");
            }
        }
        if (chunk_size == 0) {
            break;
        }
        if (safe_write(ctx->fd_out, chunk_msg, chunk_size, -1) < 0) {
            die(1, "write()");
        }
        chunk_id++;
    }
    if (readnb < 0) {
        die(1, "read()");
    }
    if (chunk_size != 0) {
        die(0, "Premature end of file");
    }
    return 0;
}

static int
read_password_file(Context *ctx, const char *file)
{
    char    password_[512], *password = password_;
    ssize_t password_len;
    int     fd;

    fd = file_open(file, 0);
    if ((password_len = safe_read(fd, password, sizeof password_)) < 0) {
        die(1, "Unable to read the password");
    }
    while (password_len > 0 &&
           (password[password_len - 1] == ' ' || password[password_len - 1] == '\r' ||
            password[password_len - 1] == '\n')) {
        password_len--;
    }
    while (password_len > 0 && (*password == ' ' || *password == '\r' || *password == '\n')) {
        password++;
        password_len--;
    }
    if (password_len <= 0) {
        die(0, "Empty password");
    }
    close(fd);
    derive_key(ctx, password, password_len);

    return 0;
}

static void
read_password_from_terminal(Context *ctx)
{
    char  password[512];
    char *p;

    printf("Password: ");
    fflush(stdout);
    if (fgets(password, sizeof password, stdin) == NULL) {
        die(1, "fgets()");
    }
    if ((p = strchr(password, '\r')) != NULL) {
        *p = 0;
    }
    if ((p = strchr(password, '\n')) != NULL) {
        *p = 0;
    }
    derive_key(ctx, password, strlen(password));
}

static void
passgen(void)
{
    unsigned char password[32];
    char          hex[32 * 2 + 1];

    randombytes(password, sizeof password);
    bin2hex(hex, sizeof hex, password, sizeof password);
    puts(hex);
    memzero(password, sizeof password);
    memzero(hex, sizeof hex);
    exit(0);
}

static void
options_parse(Context *ctx, int argc, char *argv[])
{
    int opt_flag;
    int option_index = 0;

    ctx->encrypt = -1;
    ctx->in      = NULL;
    ctx->out     = NULL;
    optind       = 0;
#ifdef _OPTRESET
    optreset = 1;
#endif
    while ((opt_flag = getopt_long(argc, argv, getopt_options, getopt_long_options,
                                   &option_index)) != -1) {
        switch (opt_flag) {
        case 'd':
            ctx->encrypt = 0;
            break;
        case 'e':
            ctx->encrypt = 1;
            break;
        case 'G':
            passgen();
            break; /* NOTREACHED */
        case 'h':
            usage(stdout);
            break; /* NOTREACHED */
        case 'i':
            ctx->in = optarg;
            break;
        case 'o':
            ctx->out = optarg;
            break;
        case 'p':
            derive_key(ctx, optarg, strlen(optarg));
            break;
        case 'P':
            read_password_file(ctx, optarg);
            break;
        default:
            usage(stderr);
        }
    }
    if (ctx->encrypt == -1) {
        usage(stderr);
    }
    if (ctx->has_key == 0) {
        read_password_from_terminal(ctx);
    }
}

int
main(int argc, char *argv[])
{
    Context ctx;

    memset(&ctx, 0, sizeof ctx);
    options_parse(&ctx, argc, argv);
    ctx.sizeof_buf = DEFAULT_BUFFER_SIZE;
    if (ctx.sizeof_buf < MIN_BUFFER_SIZE) {
        ctx.sizeof_buf = MIN_BUFFER_SIZE;
    } else if (ctx.sizeof_buf > MAX_BUFFER_SIZE) {
        ctx.sizeof_buf = MAX_BUFFER_SIZE;
    }
    if ((ctx.buf = (unsigned char *) malloc(ctx.sizeof_buf)) == NULL) {
        die(1, "malloc()");
    }

    ctx.fd_in  = file_open(ctx.in, 0);
    ctx.fd_out = file_open(ctx.out, 1);
    if (ctx.encrypt) {
        stream_encrypt(&ctx);
    } else {
        stream_decrypt(&ctx);
    }
    free(ctx.buf);
    close(ctx.fd_out);
    close(ctx.fd_in);
    memzero(&ctx, sizeof ctx);

    return 0;
}