/*
 * otp_pad.c — OTP pad state for n_signer.
 *
 * One pad per session (by design — see plans/otp_nostr_integration.md).
 * The pad is bound at startup via otp_pad_bind() and accessed by the
 * otp_encrypt / otp_decrypt dispatcher verbs via otp_pad_get_state().
 *
 * The pad file is opened read-only and kept open for the lifetime of the
 * process. The per-pad .state file (offset counter) is read and written
 * via libotppad. Offset writes are atomic (temp + rename) inside libotppad.
 *
 * Pad bytes are never loaded whole into RAM. Each request seeks to the
 * current offset and reads exactly the slice it needs into a small
 * mlock'd scratch buffer.
 */
#define _POSIX_C_SOURCE 200809L
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <dirent.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <errno.h>

#include "libotppad.h"
#include "otp_pad.h"

/* from secure_mem.h (headerless decls pattern) */
extern void secure_memzero(void *ptr, size_t len);

/* ------------------------------------------------------------------ */
/* OTP pad state                                                      */
/* ------------------------------------------------------------------ */

#define OTP_PAD_DIR_MAX     512
#define OTP_PAD_CHKSUM_MAX  128
#define OTP_PAD_PATH_MAX    (OTP_PAD_DIR_MAX + OTP_PAD_CHKSUM_MAX + 16)
#define OTP_SCRATCH_MAX     (4 * 1024 * 1024)  /* 4 MB max chunk */

typedef struct {
    int      bound;                          /* 1 if a pad is bound */
    char     pads_dir[OTP_PAD_DIR_MAX];      /* directory holding .pad/.state */
    char     chksum[OTP_PAD_CHKSUM_MAX];     /* 64-hex-char pad checksum */
    char     pad_path[OTP_PAD_PATH_MAX];     /* full path to .pad file */
    FILE    *pad_fp;                         /* read-only FILE* on .pad */
    uint64_t pad_size;                       /* total pad file size in bytes */
    int      allow_blkback;                  /* 1 if --otp-allow-blkback passed */
    /* mlock'd scratch buffer for XOR */
    void    *scratch_data;
    size_t   scratch_size;
    int      scratch_locked;
} otp_pad_state_t;

static otp_pad_state_t g_otp_pad = {0};

otp_pad_state_t *otp_pad_get_state(void) {
    return &g_otp_pad;
}

int otp_pad_is_bound(void) {
    return g_otp_pad.bound;
}

const char *otp_pad_chksum(void) {
    return g_otp_pad.bound ? g_otp_pad.chksum : NULL;
}

const char *otp_pad_dir(void) {
    return g_otp_pad.bound ? g_otp_pad.pads_dir : NULL;
}

/* ------------------------------------------------------------------ */
/* Removable-mount check (Qubes guard)                                */
/* ------------------------------------------------------------------ */

/* Return 1 if `path` lives on a blkback device (e.g. /dev/xvdi via qvm-block),
 * 0 if it's a directly-owned device (e.g. /dev/sda via PCI passthrough),
 * -1 on error. Best-effort: compares the fs source device name. */
static int is_blkback_mount(const char *path) {
    struct stat st;
    if (stat(path, &st) != 0) return -1;

    /* Read the mount source for the filesystem containing `path`. */
    FILE *mtab = fopen("/proc/mounts", "r");
    if (!mtab) return -1;

    char line[1024];
    int found = 0;
    int is_blkback = 0;
    size_t best_len = 0;

    while (fgets(line, sizeof(line), mtab)) {
        char source[512], mount[512], fstype[64];
        if (sscanf(line, "%511s %511s %63s", source, mount, fstype) < 3) continue;
        size_t mlen = strlen(mount);
        /* Pick the longest matching mount prefix. */
        if (strncmp(path, mount, mlen) == 0 &&
            (path[mlen] == '/' || path[mlen] == '\0') &&
            mlen > best_len) {
            best_len = mlen;
            found = 1;
            /* blkback devices show up as /dev/xvd* in the guest. */
            is_blkback = (strncmp(source, "/dev/xvd", 8) == 0);
        }
    }
    fclose(mtab);
    if (!found) return -1;
    return is_blkback;
}

/* ------------------------------------------------------------------ */
/* Find a pad by chksum prefix in pads_dir                            */
/* ------------------------------------------------------------------ */

/* Resolve a chksum-or-prefix to a full 64-char chksum by scanning pads_dir.
 * Returns 0 on success and fills `out_chksum` (must be >= OTP_PAD_CHKSUM_MAX).
 * Returns -1 if not found, -2 if ambiguous (multiple matches). */
static int resolve_pad_chksum(const char *pads_dir, const char *prefix,
                              char *out_chksum) {
    DIR *d = opendir(pads_dir);
    if (!d) return -1;

    struct dirent *e;
    int matches = 0;
    char found[OTPPAD_CHKSUM_HEX_LEN + 1] = {0};
    size_t plen = strlen(prefix);

    while ((e = readdir(d)) != NULL) {
        size_t nlen = strlen(e->d_name);
        if (nlen < 5 || strcmp(e->d_name + nlen - 4, ".pad") != 0) continue;
        size_t base_len = nlen - 4;  /* without ".pad" */
        if (base_len != OTPPAD_CHKSUM_HEX_LEN) continue;
        if (plen == 0 || strncmp(e->d_name, prefix, plen) == 0) {
            memcpy(found, e->d_name, base_len);
            found[base_len] = '\0';
            matches++;
        }
    }
    closedir(d);

    if (matches == 0) return -1;
    if (matches > 1) return -2;
    strncpy(out_chksum, found, OTPPAD_CHKSUM_HEX_LEN);
    out_chksum[OTPPAD_CHKSUM_HEX_LEN] = '\0';
    return 0;
}

/* ------------------------------------------------------------------ */
/* Bind / unbind                                                      */
/* ------------------------------------------------------------------ */

/* Bind a pad at startup. Returns 0 on success, non-zero on error.
 * `pads_dir` is the directory containing <chksum>.pad and <chksum>.state.
 * `pad_spec` is a full 64-char chksum or a unique prefix.
 * `allow_blkback` non-zero skips the blkback guard (for qvm-block testing). */
int otp_pad_bind(const char *pads_dir, const char *pad_spec, int allow_blkback) {
    if (!pads_dir || !pad_spec) return 1;
    if (g_otp_pad.bound) return 2;  /* already bound */

    /* Removable-mount guard. */
    int blkback = is_blkback_mount(pads_dir);
    if (blkback < 0) {
        /* Could not determine — warn but continue (best-effort). */
        fprintf(stderr, "otp_pad: warning: could not determine mount type for %s\n",
                pads_dir);
    } else if (blkback && !allow_blkback) {
        fprintf(stderr, "otp_pad: %s is on a blkback device (qvm-block).\n",
                pads_dir);
        fprintf(stderr, "  Refusing to bind for pad secrecy. Use PCI USB "
                        "controller passthrough, or pass --otp-allow-blkback "
                        "to override (not recommended for production pads).\n");
        return 3;
    }

    /* Resolve the pad chksum. */
    char chksum[OTPPAD_CHKSUM_HEX_LEN + 1];
    if (strlen(pad_spec) == OTPPAD_CHKSUM_HEX_LEN) {
        strncpy(chksum, pad_spec, OTPPAD_CHKSUM_HEX_LEN);
        chksum[OTPPAD_CHKSUM_HEX_LEN] = '\0';
        /* Verify the file exists. */
        char path[OTP_PAD_PATH_MAX];
        snprintf(path, sizeof(path), "%s/%s.pad", pads_dir, chksum);
        if (access(path, R_OK) != 0) {
            fprintf(stderr, "otp_pad: pad file not found: %s\n", path);
            return 4;
        }
    } else {
        int r = resolve_pad_chksum(pads_dir, pad_spec, chksum);
        if (r == -1) {
            fprintf(stderr, "otp_pad: no pad matching prefix '%s' in %s\n",
                    pad_spec, pads_dir);
            return 5;
        } else if (r == -2) {
            fprintf(stderr, "otp_pad: ambiguous pad prefix '%s' (multiple matches)\n",
                    pad_spec);
            return 6;
        }
    }

    /* Build the pad path and open it read-only. */
    char pad_path[OTP_PAD_PATH_MAX];
    snprintf(pad_path, sizeof(pad_path), "%s/%s.pad", pads_dir, chksum);

    FILE *fp = fopen(pad_path, "rb");
    if (!fp) {
        fprintf(stderr, "otp_pad: cannot open %s: %s\n", pad_path, strerror(errno));
        return 7;
    }

    /* Determine pad size. */
    struct stat st;
    if (fstat(fileno(fp), &st) != 0) {
        fprintf(stderr, "otp_pad: fstat failed: %s\n", strerror(errno));
        fclose(fp);
        return 8;
    }
    if (st.st_size < (off_t)OTPPAD_HEADER_RESERVED) {
        fprintf(stderr, "otp_pad: pad too small (%lld bytes, need >= %d)\n",
                (long long)st.st_size, OTPPAD_HEADER_RESERVED);
        fclose(fp);
        return 9;
    }

    /* Verify the pad checksum matches the filename. */
    char computed[OTPPAD_CHKSUM_HEX_LEN + 1];
    if (otppad_checksum(pad_path, computed) != 0) {
        fprintf(stderr, "otp_pad: checksum computation failed\n");
        fclose(fp);
        return 10;
    }
    if (strcmp(computed, chksum) != 0) {
        fprintf(stderr, "otp_pad: checksum mismatch (file says %s, computed %s)\n",
                chksum, computed);
        fclose(fp);
        return 11;
    }

    /* Read the current offset. */
    uint64_t offset;
    if (otppad_state_read(pads_dir, chksum, &offset) != 0) {
        /* No state file — initialize at the reserved header. */
        offset = OTPPAD_HEADER_RESERVED;
        if (otppad_state_write(pads_dir, chksum, offset) != 0) {
            fprintf(stderr, "otp_pad: cannot write initial state file\n");
            fclose(fp);
            return 12;
        }
    }
    if (offset < OTPPAD_HEADER_RESERVED) {
        fprintf(stderr, "otp_pad: offset %llu < reserved header %d\n",
                (unsigned long long)offset, OTPPAD_HEADER_RESERVED);
        fclose(fp);
        return 13;
    }
    if (offset > (uint64_t)st.st_size) {
        fprintf(stderr, "otp_pad: offset %llu past end of pad (%lld)\n",
                (unsigned long long)offset, (long long)st.st_size);
        fclose(fp);
        return 14;
    }

    /* Commit. */
    strncpy(g_otp_pad.pads_dir, pads_dir, OTP_PAD_DIR_MAX - 1);
    strncpy(g_otp_pad.chksum, chksum, OTP_PAD_CHKSUM_MAX - 1);
    strncpy(g_otp_pad.pad_path, pad_path, OTP_PAD_PATH_MAX - 1);
    g_otp_pad.pad_fp = fp;
    g_otp_pad.pad_size = (uint64_t)st.st_size;
    g_otp_pad.allow_blkback = allow_blkback;
    g_otp_pad.bound = 1;

    fprintf(stderr, "otp_pad: bound pad %s (%llu bytes, offset=%llu)\n",
            chksum, (unsigned long long)st.st_size,
            (unsigned long long)offset);
    return 0;
}

void otp_pad_unbind(void) {
    if (g_otp_pad.pad_fp) {
        fclose(g_otp_pad.pad_fp);
        g_otp_pad.pad_fp = NULL;
    }
    if (g_otp_pad.scratch_data) {
        secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        if (g_otp_pad.scratch_locked) {
            munlock(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        }
        free(g_otp_pad.scratch_data);
        g_otp_pad.scratch_data = NULL;
        g_otp_pad.scratch_size = 0;
        g_otp_pad.scratch_locked = 0;
    }
    secure_memzero(&g_otp_pad, sizeof(g_otp_pad));
}

/* ------------------------------------------------------------------ */
/* Core encrypt/decrypt transform                                     */
/* ------------------------------------------------------------------ */

/* Ensure the scratch buffer is at least `size` bytes, mlock'd. */
static int ensure_scratch(size_t size) {
    if (size > OTP_SCRATCH_MAX) return -1;
    if (g_otp_pad.scratch_size >= size && g_otp_pad.scratch_data) return 0;
    /* Grow. */
    if (g_otp_pad.scratch_data) {
        secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        if (g_otp_pad.scratch_locked) {
            munlock(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        }
        free(g_otp_pad.scratch_data);
        g_otp_pad.scratch_data = NULL;
        g_otp_pad.scratch_size = 0;
        g_otp_pad.scratch_locked = 0;
    }
    g_otp_pad.scratch_data = malloc(size);
    if (!g_otp_pad.scratch_data) return -2;
    g_otp_pad.scratch_size = size;
    if (mlock(g_otp_pad.scratch_data, size) == 0) {
        g_otp_pad.scratch_locked = 1;
    }
    return 0;
}

/* Read `len` bytes from the pad at `offset` into `out` (must be mlock'd by
 * caller). Returns 0 on success, non-zero on error. */
static int read_pad_slice(uint64_t offset, size_t len, unsigned char *out) {
    if (!g_otp_pad.pad_fp) return -1;
    if (fseek(g_otp_pad.pad_fp, (long)offset, SEEK_SET) != 0) return -2;
    size_t got = fread(out, 1, len, g_otp_pad.pad_fp);
    if (got != len) return -3;
    return 0;
}

/* ------------------------------------------------------------------ */
/* Public encrypt/decrypt entrypoints (called by dispatcher verbs)   */
/* ------------------------------------------------------------------ */

/* OTPPAD_ENCRYPT_OK etc. are returned via *out_result. */

/* Encrypt: takes plaintext bytes, returns malloc'd ASCII armor or binary blob.
 *
 * `encoding` is "ascii" or "binary".
 * On success returns 0 and sets *out_payload (malloc'd, caller frees),
 * *out_payload_len, and *out_new_offset.
 */
int otp_pad_encrypt(const unsigned char *plaintext, size_t pt_len,
                    const char *encoding,
                    char **out_payload, size_t *out_payload_len,
                    uint64_t *out_new_offset) {
    if (!g_otp_pad.bound) return 1;  /* not bound */
    if (!plaintext || !out_payload || !out_payload_len || !out_new_offset) return 2;
    *out_payload = NULL;
    *out_payload_len = 0;

    /* Pad the plaintext. */
    size_t chunk = otppad_chunk_size(pt_len);
    if (ensure_scratch(chunk) != 0) return 3;
    unsigned char *buf = (unsigned char *)g_otp_pad.scratch_data;
    memcpy(buf, plaintext, pt_len);
    if (otppad_pad_apply(buf, pt_len, chunk) != 0) return 4;

    /* Read current offset. */
    uint64_t offset;
    if (otppad_state_read(g_otp_pad.pads_dir, g_otp_pad.chksum, &offset) != 0) {
        return 5;
    }
    if (offset + chunk > g_otp_pad.pad_size) {
        return 6;  /* pad exhausted */
    }

    /* Read the pad slice into a second scratch buffer. */
    /* Reuse the same scratch: read pad into second half, XOR in place.
     * For simplicity, allocate a separate pad-slice buffer. */
    unsigned char *pad_slice = (unsigned char *)malloc(chunk);
    if (!pad_slice) return 7;
    if (read_pad_slice(offset, chunk, pad_slice) != 0) {
        free(pad_slice);
        return 8;
    }

    /* XOR. */
    for (size_t i = 0; i < chunk; i++) {
        buf[i] ^= pad_slice[i];
    }
    secure_memzero(pad_slice, chunk);
    free(pad_slice);

    /* Advance offset atomically. */
    uint64_t new_offset = offset + chunk;
    if (otppad_state_write(g_otp_pad.pads_dir, g_otp_pad.chksum, new_offset) != 0) {
        return 9;
    }

    /* Encode output. */
    if (encoding && strcmp(encoding, "binary") == 0) {
        /* Binary .otp: header + encrypted (padded) data. */
        otppad_bin_header_t hdr;
        memset(&hdr, 0, sizeof(hdr));
        memcpy(hdr.magic, OTPPAD_MAGIC, OTPPAD_MAGIC_LEN);
        hdr.version = OTPPAD_FORMAT_VERSION;
        /* pad_chksum is binary 32 bytes — convert hex to bytes. */
        for (int i = 0; i < OTPPAD_CHKSUM_BIN_LEN; i++) {
            unsigned int byte;
            sscanf(g_otp_pad.chksum + i * 2, "%02x", &byte);
            hdr.pad_chksum[i] = (unsigned char)byte;
        }
        hdr.pad_offset = offset;
        hdr.file_mode = 0644;
        hdr.file_size = pt_len;  /* original (unpadded) size */

        /* Build the blob in memory (avoid fmemopen — it can misbehave
         * with NUL bytes on some platforms). */
        size_t blob_size = 58 + chunk;
        unsigned char *blob = (unsigned char *)malloc(blob_size);
        if (!blob) return 10;
        unsigned char *p = blob;
        memcpy(p, OTPPAD_MAGIC, 4); p += 4;
        memcpy(p, &hdr.version, 2); p += 2;
        memcpy(p, hdr.pad_chksum, OTPPAD_CHKSUM_BIN_LEN); p += OTPPAD_CHKSUM_BIN_LEN;
        memcpy(p, &hdr.pad_offset, 8); p += 8;
        memcpy(p, &hdr.file_mode, 4); p += 4;
        memcpy(p, &hdr.file_size, 8); p += 8;
        /* p is now at byte 58. Copy the encrypted (padded) data. */
        memcpy(p, buf, chunk);
        *out_payload = (char *)blob;
        *out_payload_len = blob_size;
    } else {
        /* ASCII armor. */
        char *armor = NULL;
        if (otppad_armor_generate(NSIGNER_OTP_VERSION, g_otp_pad.chksum, offset,
                                  buf, chunk, &armor) != 0) {
            return 14;
        }
        *out_payload = armor;
        *out_payload_len = strlen(armor);
    }

    *out_new_offset = new_offset;
    secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
    return 0;
}

/* Decrypt: takes ASCII armor or binary blob, returns malloc'd plaintext.
 *
 * `encoding` is "ascii" or "binary" (auto-detected if NULL).
 * On success returns 0 and sets *out_plaintext (malloc'd, caller frees),
 * *out_pt_len.
 */
int otp_pad_decrypt(const char *input, size_t input_len,
                    const char *encoding,
                    unsigned char **out_plaintext, size_t *out_pt_len) {
    if (!g_otp_pad.bound) return 1;
    if (!input || !out_plaintext || !out_pt_len) return 2;
    *out_plaintext = NULL;
    *out_pt_len = 0;

    uint64_t offset;
    size_t chunk;
    unsigned char *ciphertext = NULL;
    size_t ct_len = 0;

    int is_binary;
    if (encoding && strcmp(encoding, "binary") == 0) {
        is_binary = 1;
    } else if (encoding && strcmp(encoding, "ascii") == 0) {
        is_binary = 0;
    } else {
        /* Auto-detect by magic bytes. */
        is_binary = (input_len >= 4 && memcmp(input, OTPPAD_MAGIC, 4) == 0);
    }

    if (!is_binary) {
        /* ASCII armor. */
        char chksum[OTPPAD_CHKSUM_HEX_LEN + 1];
        char b64[65536];
        if (otppad_armor_parse(input, chksum, &offset, b64, sizeof(b64)) != 0) {
            return 3;
        }
        if (strcmp(chksum, g_otp_pad.chksum) != 0) {
            return 4;  /* pad mismatch */
        }
        int dlen = 0;
        ciphertext = otppad_base64_decode(b64, &dlen);
        if (!ciphertext) return 5;
        ct_len = (size_t)dlen;
        chunk = ct_len;
    } else {
        /* Binary .otp blob — parse directly from the buffer (avoid fmemopen
         * which can misbehave with NUL bytes on some platforms). */
        if (input_len < 58) return 6;
        const unsigned char *p = (const unsigned char *)input;
        otppad_bin_header_t hdr;
        memset(&hdr, 0, sizeof(hdr));
        memcpy(hdr.magic, p, 4); p += 4;
        memcpy(&hdr.version, p, 2); p += 2;
        memcpy(hdr.pad_chksum, p, OTPPAD_CHKSUM_BIN_LEN); p += OTPPAD_CHKSUM_BIN_LEN;
        memcpy(&hdr.pad_offset, p, 8); p += 8;
        memcpy(&hdr.file_mode, p, 4); p += 4;
        memcpy(&hdr.file_size, p, 8); p += 8;
        /* p is now at byte 58. */
        if (memcmp(hdr.magic, OTPPAD_MAGIC, OTPPAD_MAGIC_LEN) != 0) {
            return 8;
        }
        /* Convert binary chksum to hex for comparison. */
        char chksum_hex[OTPPAD_CHKSUM_HEX_LEN + 1];
        for (int i = 0; i < OTPPAD_CHKSUM_BIN_LEN; i++) {
            sprintf(chksum_hex + i * 2, "%02x", hdr.pad_chksum[i]);
        }
        chksum_hex[OTPPAD_CHKSUM_HEX_LEN] = '\0';
        if (strcmp(chksum_hex, g_otp_pad.chksum) != 0) {
            return 9;
        }
        offset = hdr.pad_offset;
        /* Remaining bytes after the 58-byte header are the ciphertext. */
        ct_len = input_len - 58;
        chunk = ct_len;
        ciphertext = (unsigned char *)malloc(ct_len);
        if (!ciphertext) return 10;
        memcpy(ciphertext, p, ct_len);
    }

    if (ensure_scratch(chunk) != 0) {
        free(ciphertext); return 12;
    }
    unsigned char *buf = (unsigned char *)g_otp_pad.scratch_data;

    /* Read pad slice. */
    unsigned char *pad_slice = (unsigned char *)malloc(chunk);
    if (!pad_slice) { free(ciphertext); return 13; }
    if (read_pad_slice(offset, chunk, pad_slice) != 0) {
        free(pad_slice); free(ciphertext); return 14;
    }

    /* XOR (in-place into scratch). */
    for (size_t i = 0; i < chunk; i++) {
        buf[i] = ciphertext[i] ^ pad_slice[i];
    }
    secure_memzero(pad_slice, chunk);
    free(pad_slice);
    secure_memzero(ciphertext, ct_len);
    free(ciphertext);

    /* Strip padding. */
    size_t pt_len;
    if (otppad_pad_remove(buf, chunk, &pt_len) != 0) {
        secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        return 15;
    }

    /* Return plaintext. */
    unsigned char *pt = (unsigned char *)malloc(pt_len ? pt_len : 1);
    if (!pt) {
        secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);
        return 16;
    }
    memcpy(pt, buf, pt_len);
    secure_memzero(g_otp_pad.scratch_data, g_otp_pad.scratch_size);

    *out_plaintext = pt;
    *out_pt_len = pt_len;
    return 0;
}

/* ------------------------------------------------------------------ */
/* Helpers for the dispatcher                                         */
/* ------------------------------------------------------------------ */

uint64_t otp_pad_current_offset(void) {
    if (!g_otp_pad.bound) return 0;
    uint64_t off;
    if (otppad_state_read(g_otp_pad.pads_dir, g_otp_pad.chksum, &off) != 0) {
        return 0;
    }
    return off;
}

uint64_t otp_pad_size(void) {
    return g_otp_pad.bound ? g_otp_pad.pad_size : 0;
}
