#include <stdio.h>
#include <libelf.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <errno.h>
#include <iostream>

enum LibraryType {
    LIBRARY_TYPE_ERROR,
    LIBRARY_TYPE_32_BIT,
    LIBRARY_TYPE_64_BIT,
};

LibraryType TestLibrary(const char * libraryName) {
        // Initialize libelf - must be called before any other libelf function
    if (elf_version(EV_CURRENT) == EV_NONE) {
        std::cerr << "Failed to initialize libelf: " << elf_errmsg(-1) << std::endl;
        return LIBRARY_TYPE_ERROR;
    }

    // Check whether the library file class is 32-bit or 64-bit 
    int fd = open(libraryName, O_RDONLY);
    if (fd == -1)
    {
        std::cerr << "Failed to open file: " << strerror(errno) << std::endl;
        return LIBRARY_TYPE_ERROR;
    }

    Elf* elf = elf_begin(fd, ELF_C_READ, NULL);
    if (elf == NULL)
    {
        std::cerr << "Failed to initialize ELF: " << elf_errmsg(-1) << std::endl;
        close(fd);
        return LIBRARY_TYPE_ERROR;
    }

    Elf32_Ehdr* ehdr32 = elf32_getehdr(elf);
    Elf64_Ehdr* ehdr64 = elf64_getehdr(elf);

    if (ehdr32 != NULL)
    {
        std::cout << "Library file is 32-bit" << std::endl;
        return LibraryType::LIBRARY_TYPE_32_BIT;
    }
    else if (ehdr64 != NULL)
    {
        std::cout << "Library file is 64-bit" << std::endl;
        return LibraryType::LIBRARY_TYPE_64_BIT;
    }

    std::cerr << "Failed to retrieve ELF header" << std::endl;


    elf_end(elf);
    close(fd);
    return LibraryType::LIBRARY_TYPE_ERROR;
}

bool FindTargetSection32(Elf_Scn *&scn, Elf *&elf, Elf32_Shdr *&shdr, const char *targetSectionName, size_t shstrndx) {
    char *curSectionName = NULL;
    while ((scn = elf_nextscn(elf, scn)) != NULL) {
        // Retrieve the section header for 32-bit ELF files
        shdr = elf32_getshdr(scn);
        if (shdr == NULL) {
            fprintf(stderr, "Error: elf32_getshdr failed: %s\n", elf_errmsg(-1));
            scn = nullptr;
            return false;
        }

        // Get the name of the section
        curSectionName = elf_strptr(elf, shstrndx, shdr->sh_name);
        if (curSectionName == NULL) {
            fprintf(stderr, "Error: elf_strptr failed: %s\n", elf_errmsg(-1));
            scn = nullptr;
            return false;
        }

        // Check if the section is .note.gnu.build-id
        if (strcmp(curSectionName, targetSectionName) == 0) {
            return true;
        }
    }
    return false;
}

bool FindTargetSection64(Elf_Scn *&scn, Elf *&elf, Elf64_Shdr *&shdr, const char *targetSectionName, size_t shstrndx) {
    char *curSectionName = NULL;
    while ((scn = elf_nextscn(elf, scn)) != NULL) {
        // Retrieve the section header for 64-bit ELF files
        shdr = elf64_getshdr(scn);
        if (shdr == NULL) {
            fprintf(stderr, "Error: elf64_getshdr failed: %s\n", elf_errmsg(-1));
            scn = nullptr;
            return false;
        }

        // Get the name of the section
        curSectionName = elf_strptr(elf, shstrndx, shdr->sh_name);
        if (curSectionName == NULL) {
            fprintf(stderr, "Error: elf_strptr failed: %s\n", elf_errmsg(-1));
            scn = nullptr;
            return false;
        }

        // Check if the section is .note.gnu.build-id
        if (strcmp(curSectionName, targetSectionName) == 0) {
            return true;
        }
    }
    return false;
}

// extract build id section and build id, store the id as output file name. 
// extract the trace dict section, store as output file content. 
int ExtractSectionFromLibraryAndSave32bit(const char * libraryName, const char * targetSectionName){
    static const char * const BUILD_ID_SECTION_NAME = ".note.gnu.build-id";
    // static const char * const TRACE_DICT_SECTION_NAME = ".rodata";
    // static const char * const OUTPUT_FILE_NAME = "data_section.txt";
    Elf *elf = NULL;
    Elf_Scn *scn = NULL;
    Elf_Data *data = NULL;
    Elf32_Shdr *shdr = NULL;  // Use Elf32_Shdr for 32-bit ELF files
    char *name = NULL;
    int fd_in = 0;
    int fd_out = 0;
    size_t shstrndx = 0;

    // Initialize libelf library
    if (elf_version(EV_CURRENT) == EV_NONE) {
        fprintf(stderr, "Error: ELF library initialization failed: %s\n", elf_errmsg(-1));
        return EXIT_FAILURE;
    }

    // Open the input ELF file
    fd_in = open(libraryName, O_RDONLY, 0);
    if (fd_in < 0) {
        perror("Error opening input ELF file");
        return EXIT_FAILURE;
    }

    // Convert the file descriptor to an ELF handle
    elf = elf_begin(fd_in, ELF_C_READ, NULL);
    if (elf == NULL) {
        fprintf(stderr, "Error: elf_begin failed: %s\n", elf_errmsg(-1));
        close(fd_in);
        return EXIT_FAILURE;
    }

    // Retrieve the index of the section name string table
    if (elf_getshdrstrndx(elf, &shstrndx) != 0) {
        fprintf(stderr, "Error: elf_getshdrstrndx failed: %s\n", elf_errmsg(-1));
        goto failure;
    }

    if (!FindTargetSection32(scn, elf, shdr, BUILD_ID_SECTION_NAME, shstrndx)) {
        fprintf(stderr, "Error: Failed to find target section: %s\n", BUILD_ID_SECTION_NAME);
        goto failure;
    }

    // get build id content into array : get the 0x10 to 0x20 bytes of its content : 
    data = elf_getdata(scn, data);
    if (data == NULL) {
        fprintf(stderr, "Error: get build id failed: %s\n", elf_errmsg(-1));
        goto failure;
    }
    if (data->d_size < 0x20) {
        fprintf(stderr, "Error: build id section size is less than 0x20\n");
        goto failure;
    }
    unsigned char buildIdBuffer[20]; // 20 bytes for build id, need only store 16 bytes
    memcpy(buildIdBuffer, static_cast<char*>(data->d_buf) + 0x10, 0x10);
    // print build id byte by byte in hex : 
    printf("build id is: ");
    for (int i = 0; i < 0x10; i++) {
        printf("%02x ", buildIdBuffer[i]);
    }
    printf("\n");
    char buildIdHexName[40];
    // convert build id to hex string
    for (int i = 0; i < 0x10; i++) {
        sprintf(buildIdHexName + i * 2, "%02x", buildIdBuffer[i]);
    }
    buildIdHexName[32] = '\0';
    // printf("%s\n", buildIdHexName);
    
    scn = nullptr;
    if (!FindTargetSection32(scn, elf, shdr, targetSectionName, shstrndx)) {
        printf("Error: Failed to find target section: %s\n", targetSectionName);
        goto failure;
    }

    fd_out = open(buildIdHexName, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
    if (fd_out < 0) {
        printf("Error opening output file");
        goto failure;
    }

    data = nullptr;
    // Read the section and write it to the file
    while ((data = elf_getdata(scn, data)) != NULL) {
        if (write(fd_out, data->d_buf, data->d_size) != (ssize_t)data->d_size) {
            printf("Error writing to output file");
            close(fd_out);
            goto failure;
        }
    }

    // Close the output file descriptor
    close(fd_out);

success:
    elf_end(elf);
    close(fd_in);
    return EXIT_SUCCESS;

failure:
    elf_end(elf);
    close(fd_in);
    return EXIT_FAILURE;
}


// extract build id section and build id, store the id as output file name. 
// extract the trace dict section, store as output file content. 
int ExtractSectionFromLibraryAndSave64bit(const char * libraryName, const char * targetSectionName){
    static const char * const BUILD_ID_SECTION_NAME = ".note.gnu.build-id";
    // static const char * const TRACE_DICT_SECTION_NAME = ".rodata";
    // static const char * const OUTPUT_FILE_NAME = "data_section.txt";
    Elf *elf = NULL;
    Elf_Scn *scn = NULL;
    Elf_Data *data = NULL;
    Elf64_Shdr *shdr = NULL;  // Use Elf64_Shdr for 64-bit ELF files
    char *name = NULL;
    int fd_in = 0;
    int fd_out = 0;
    size_t shstrndx = 0;

    // Initialize libelf library
    if (elf_version(EV_CURRENT) == EV_NONE) {
        fprintf(stderr, "Error: ELF library initialization failed: %s\n", elf_errmsg(-1));
        return EXIT_FAILURE;
    }

    // Open the input ELF file
    fd_in = open(libraryName, O_RDONLY, 0);
    if (fd_in < 0) {
        perror("Error opening input ELF file");
        return EXIT_FAILURE;
    }

    // Convert the file descriptor to an ELF handle
    elf = elf_begin(fd_in, ELF_C_READ, NULL);
    if (elf == NULL) {
        fprintf(stderr, "Error: elf_begin failed: %s\n", elf_errmsg(-1));
        close(fd_in);
        return EXIT_FAILURE;
    }

    // Retrieve the index of the section name string table
    if (elf_getshdrstrndx(elf, &shstrndx) != 0) {
        fprintf(stderr, "Error: elf_getshdrstrndx failed: %s\n", elf_errmsg(-1));
        goto failure;
    }

    if (!FindTargetSection64(scn, elf, shdr, BUILD_ID_SECTION_NAME, shstrndx)) {
        fprintf(stderr, "Error: Failed to find target section: %s\n", BUILD_ID_SECTION_NAME);
        goto failure;
    }

    // get build id content into array : get the 0x10 to 0x20 bytes of its content : 
    data = elf_getdata(scn, data);
    if (data == NULL) {
        fprintf(stderr, "Error: get build id failed: %s\n", elf_errmsg(-1));
        goto failure;
    }
    if (data->d_size < 0x20) {
        fprintf(stderr, "Error: build id section size is less than 0x20\n");
        goto failure;
    }
    unsigned char buildIdBuffer[20]; // 20 bytes for build id, need only store 16 bytes
    memcpy(buildIdBuffer, static_cast<char*>(data->d_buf) + 0x10, 0x10);
    // print build id byte by byte in hex : 
    printf("build id is: ");
    for (int i = 0; i < 0x10; i++) {
        printf("%02x ", buildIdBuffer[i]);
    }
    printf("\n");
    char buildIdHexName[40];
    // convert build id to hex string
    for (int i = 0; i < 0x10; i++) {
        sprintf(buildIdHexName + i * 2, "%02x", buildIdBuffer[i]);
    }
    buildIdHexName[32] = '\0';
    // printf("%s\n", buildIdHexName);

    scn = nullptr;
    if (!FindTargetSection64(scn, elf, shdr, targetSectionName, shstrndx)) {
        printf("Error: Failed to find target section: %s\n", targetSectionName);
        goto failure;
    }

    fd_out = open(buildIdHexName, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
    if (fd_out < 0) {
        printf("Error opening output file");
        goto failure;
    }

    data = nullptr;
    // Read the section and write it to the file
    while ((data = elf_getdata(scn, data)) != NULL) {
        if (write(fd_out, data->d_buf, data->d_size) != (ssize_t)data->d_size) {
            printf("Error writing to output file");
            close(fd_out);
            goto failure;
        }
    }

    // Close the output file descriptor
    close(fd_out);

success:
    elf_end(elf);
    close(fd_in);
    return EXIT_SUCCESS;

failure:
    elf_end(elf);
    close(fd_in);
    return EXIT_FAILURE;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <library_name> <section_name> \n", argv[0]);
        return EXIT_FAILURE;
    }
    LibraryType libraryType = TestLibrary(argv[1]);
    if (libraryType == LibraryType::LIBRARY_TYPE_32_BIT)
        return ExtractSectionFromLibraryAndSave32bit(argv[1], argv[2]);
    else if (libraryType == LibraryType::LIBRARY_TYPE_64_BIT)
        return ExtractSectionFromLibraryAndSave64bit(argv[1], argv[2]);

    return EXIT_FAILURE;
}