main.c

#ifndef _WIN32
#define _POSIX_C_SOURCE 200809L
#endif
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <limits.h>
#include <unistd.h>
#include <getopt.h>

#include "compat.h"
#include "crypto.h"
#include "password.h"

static int verbose_flag = 0;

/* -------------------------------------------------------------------------
 * I/O helpers
 * ---------------------------------------------------------------------- */

static int copy_to_stdout(const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "error: cannot open %s\n", path); return -1; }
    char buf[65536];
    size_t n;
    int write_err = 0;
    while ((n = fread(buf, 1, sizeof(buf), f)) > 0) {
        if (fwrite(buf, 1, n, stdout) != n) { write_err = 1; break; }
    }
    int err = ferror(f) || write_err;
    fclose(f);
    return err ? -1 : 0;
}

/* Copy stdin to a new temp file.  path must be at least 256 bytes. */
static int stdin_to_temp(char *path, size_t path_len) {
    int fd = portable_mkstemp(path, path_len, "in");
    if (fd < 0) { perror("mkstemp"); return -1; }
    FILE *f = fdopen(fd, "wb");
    if (!f) { close(fd); remove(path); return -1; }
    char buf[65536];
    size_t n;
    while ((n = fread(buf, 1, sizeof(buf), stdin)) > 0)
        fwrite(buf, 1, n, f);
    int err = ferror(stdin) || ferror(f);
    fclose(f);
    if (err) { remove(path); return -1; }
    return 0;
}

/* Create an empty temp file.  path must be at least 256 bytes. */
static int new_temp(char *path, size_t path_len) {
    int fd = portable_mkstemp(path, path_len, "out");
    if (fd < 0) { perror("mkstemp"); return -1; }
    close(fd);
    return 0;
}

/* -------------------------------------------------------------------------
 * Operation dispatcher
 * ---------------------------------------------------------------------- */

typedef struct {
    const uint8_t *password;
    size_t         pass_len;
    int            chunk_size;
    int            mode;  /* 0 = encrypt, 1 = decrypt */
} OpCtx;

static int run_op(const OpCtx *ctx, const char *src, const char *dst) {
    if (verbose_flag)
        fprintf(stderr, "level=INFO msg=%s src=%s dst=%s\n",
                ctx->mode ? "decrypting" : "encrypting", src, dst);

    if (ctx->mode == 0)
        return encrypt_file(src, dst, ctx->password, ctx->pass_len, ctx->chunk_size);
    else
        return decrypt_file(src, dst, ctx->password, ctx->pass_len);
}

/* Run fn once, or src→mid→dst when double is set. */
static int apply_twice(int dbl, const OpCtx *ctx, const char *src, const char *dst) {
    if (!dbl)
        return run_op(ctx, src, dst);

    char mid[256];
    if (new_temp(mid, sizeof(mid)) != 0) return -1;
    int ret = run_op(ctx, src, mid);
    if (ret == 0) ret = run_op(ctx, mid, dst);
    remove(mid);
    return ret;
}

/* -------------------------------------------------------------------------
 * Run variants
 * ---------------------------------------------------------------------- */

static int run_from_file(const char *input, const char *output_flag,
                          int stdout_flag, int double_flag, const OpCtx *ctx) {
    if (stdout_flag) {
        char dst_tmp[256];
        if (new_temp(dst_tmp, sizeof(dst_tmp)) != 0) return 1;
        int ret = apply_twice(double_flag, ctx, input, dst_tmp);
        if (ret == 0) ret = copy_to_stdout(dst_tmp) != 0 ? 1 : 0;
        remove(dst_tmp);
        return ret;
    }

    char output_buf[4096];
    const char *output;

    if (output_flag) {
        output = output_flag;
    } else if (ctx->mode == 1) {
        /* Decrypt: strip .enc or append .dec */
        size_t len = strlen(input);
        const char *suf = ".enc";
        size_t suf_len = strlen(suf);
        if (len > suf_len && strcmp(input + len - suf_len, suf) == 0) {
            memcpy(output_buf, input, len - suf_len);
            output_buf[len - suf_len] = '\0';
        } else {
            snprintf(output_buf, sizeof(output_buf), "%s.dec", input);
        }
        output = output_buf;
    } else {
        /* Encrypt: append .enc */
        snprintf(output_buf, sizeof(output_buf), "%s.enc", input);
        output = output_buf;
    }

    int ret = apply_twice(double_flag, ctx, input, output);
    if (ret == 0 && verbose_flag)
        fprintf(stderr, "level=INFO msg=%s input=%s output=%s\n",
                ctx->mode ? "decrypted" : "encrypted", input, output);
    return ret;
}

static int run_from_stdin(const char *output_flag, int double_flag, const OpCtx *ctx) {
    char src_tmp[256];
    if (stdin_to_temp(src_tmp, sizeof(src_tmp)) != 0) return 1;

    int ret = 0;
    if (output_flag) {
        ret = apply_twice(double_flag, ctx, src_tmp, output_flag) != 0 ? 1 : 0;
        if (ret == 0 && verbose_flag)
            fprintf(stderr, "level=INFO msg=%s input=stdin output=%s\n",
                    ctx->mode ? "decrypted" : "encrypted", output_flag);
    } else {
        char dst_tmp[256];
        if (new_temp(dst_tmp, sizeof(dst_tmp)) != 0) { remove(src_tmp); return 1; }
        ret = apply_twice(double_flag, ctx, src_tmp, dst_tmp) != 0 ? 1 : 0;
        if (ret == 0) ret = copy_to_stdout(dst_tmp) != 0 ? 1 : 0;
        remove(dst_tmp);
        if (ret == 0 && verbose_flag)
            fprintf(stderr, "level=INFO msg=%s input=stdin output=stdout\n",
                    ctx->mode ? "decrypted" : "encrypted");
    }

    remove(src_tmp);
    return ret;
}

/* -------------------------------------------------------------------------
 * Entry point
 * ---------------------------------------------------------------------- */

static void usage(const char *prog) {
    fprintf(stderr,
        "Usage: %s [options] [file]\n"
        "Password-based file encryption\n\n"
        "Options:\n"
        "  -d, --decrypt        Decrypt instead of encrypt\n"
        "  -o, --output FILE    Output file path\n"
        "  -p, --password PWD   Password (omit to be prompted)\n"
        "      --stdout         Write output to stdout\n"
        "  -v, --verbose        Enable verbose output\n"
        "      --chunk-size N   Encryption chunk size in bytes (default: %d)\n"
        "  -h, --help           Show this help\n",
        prog, DEFAULT_CHUNK_SIZE);
}

int main(int argc, char **argv) {
    int decrypt_flag = 0;
    int stdout_flag  = 0;
    int double_flag  = 0;
    int chunk_size   = DEFAULT_CHUNK_SIZE;
    const char *output_flag   = NULL;
    const char *password_flag = NULL;

    static struct option long_opts[] = {
        { "decrypt",    no_argument,       NULL, 'd' },
        { "output",     required_argument, NULL, 'o' },
        { "password",   required_argument, NULL, 'p' },
        { "stdout",     no_argument,       NULL, 'S' },
        { "verbose",    no_argument,       NULL, 'v' },
        { "double",     no_argument,       NULL, 'D' },  /* hidden */
        { "chunk-size", required_argument, NULL, 'C' },
        { "help",       no_argument,       NULL, 'h' },
        { NULL, 0, NULL, 0 }
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "do:p:vh", long_opts, NULL)) != -1) {
        switch (opt) {
        case 'd': decrypt_flag = 1;          break;
        case 'o': output_flag  = optarg;     break;
        case 'p': password_flag = optarg;    break;
        case 'S': stdout_flag  = 1;          break;
        case 'v': verbose_flag = 1;          break;
        case 'D': double_flag  = 1;          break;
        case 'C': {
            char *end;
            errno = 0;
            long v = strtol(optarg, &end, 10);
            if (errno || *end != '\0' || v <= 0 || v > (long)INT_MAX) {
                fprintf(stderr, "error: invalid chunk-size: %s\n", optarg);
                return 1;
            }
            chunk_size = (int)v;
            break;
        }
        case 'h': usage(argv[0]); return 0;
        default:  usage(argv[0]); return 1;
        }
    }

    if (chunk_size <= 0) {
        fprintf(stderr, "error: chunk-size must be positive\n");
        return 1;
    }

    /* Obtain password */
    uint8_t *pw;
    size_t pw_len;
    int pw_allocated = 0;

    if (password_flag) {
        pw_len = strlen(password_flag);
        pw = (uint8_t *)password_flag;
    } else {
        pw = read_password("Enter password: ", &pw_len);
        if (!pw) return 1;
        pw_allocated = 1;
    }

    if (pw_len == 0) {
        fprintf(stderr, "error: password cannot be empty\n");
        if (pw_allocated) free(pw);
        return 1;
    }

    OpCtx ctx = { pw, pw_len, chunk_size, decrypt_flag };
    int ret;

    if (optind >= argc) {
        ret = run_from_stdin(output_flag, double_flag, &ctx);
    } else {
        ret = run_from_file(argv[optind], output_flag, stdout_flag, double_flag, &ctx);
    }

    if (pw_allocated) {
        memset(pw, 0, pw_len);
        free(pw);
    }
    return ret;
}