/*
 * nso.c
 *
 * Copyright (c) 2020-2022, DarkMatterCore <pabloacurielz@gmail.com>.
 *
 * This file is part of nxdumptool (https://github.com/DarkMatterCore/nxdumptool).
 *
 * nxdumptool is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * nxdumptool is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "nxdt_utils.h"
#include "nso.h"

/* Function prototypes. */

static bool nsoGetModuleName(NsoContext *nso_ctx);
static u8 *nsoGetRodataSegment(NsoContext *nso_ctx);
static bool nsoGetModuleInfoName(NsoContext *nso_ctx, u8 *rodata_buf);
static bool nsoGetSectionFromRodataSegment(NsoContext *nso_ctx, u8 *rodata_buf, u8 **section_ptr, u64 section_offset, u64 section_size);

bool nsoInitializeContext(NsoContext *out, PartitionFileSystemContext *pfs_ctx, PartitionFileSystemEntry *pfs_entry)
{
    NcaContext *nca_ctx = NULL;
    u8 *rodata_buf = NULL;
    bool success = false, dump_nso_header = false;
    
    if (!out || !pfs_ctx || !ncaStorageIsValidContext(&(pfs_ctx->storage_ctx)) || !(nca_ctx = (NcaContext*)pfs_ctx->storage_ctx.nca_fs_ctx->nca_ctx) || \
        nca_ctx->content_type != NcmContentType_Program || !pfs_ctx->offset || !pfs_ctx->size || !pfs_ctx->is_exefs || \
        pfs_ctx->header_size <= sizeof(PartitionFileSystemHeader) || !pfs_ctx->header || !pfs_entry)
    {
        LOG_MSG("Invalid parameters!");
        return false;
    }
    
    /* Free output context beforehand. */
    nsoFreeContext(out);
    
    /* Update output context. */
    out->pfs_ctx = pfs_ctx;
    out->pfs_entry = pfs_entry;
    
    /* Get entry filename. */
    if (!(out->nso_filename = pfsGetEntryName(pfs_ctx, pfs_entry)) || !*(out->nso_filename))
    {
        LOG_MSG("Invalid Partition FS entry filename!");
        goto end;
    }
    
    /* Read NSO header. */
    if (!pfsReadEntryData(pfs_ctx, pfs_entry, &(out->nso_header), sizeof(NsoHeader), 0))
    {
        LOG_MSG("Failed to read NSO \"%s\" header!", out->nso_filename);;
        goto end;
    }
    
    /* Verify NSO header. */
    if (__builtin_bswap32(out->nso_header.magic) != NSO_HEADER_MAGIC)
    {
        LOG_MSG("Invalid NSO \"%s\" header magic word! (0x%08X != 0x%08X).", out->nso_filename, __builtin_bswap32(out->nso_header.magic), __builtin_bswap32(NSO_HEADER_MAGIC));
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.text_segment_info.file_offset < sizeof(NsoHeader) || !out->nso_header.text_segment_info.size || \
        ((out->nso_header.flags & NsoFlags_TextCompress) && (!out->nso_header.text_file_size || out->nso_header.text_file_size > out->nso_header.text_segment_info.size)) || \
        (!(out->nso_header.flags & NsoFlags_TextCompress) && out->nso_header.text_file_size != out->nso_header.text_segment_info.size) || \
        (out->nso_header.text_segment_info.file_offset + out->nso_header.text_file_size) > pfs_entry->size)
    {
        LOG_MSG("Invalid .text segment offset/size for NSO \"%s\"! (0x%08X, 0x%08X, 0x%08X).", out->nso_filename, out->nso_header.text_segment_info.file_offset, out->nso_header.text_file_size, \
                out->nso_header.text_segment_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.rodata_segment_info.file_offset < sizeof(NsoHeader) || !out->nso_header.rodata_segment_info.size || \
        ((out->nso_header.flags & NsoFlags_RoCompress) && (!out->nso_header.rodata_file_size || out->nso_header.rodata_file_size > out->nso_header.rodata_segment_info.size)) || \
        (!(out->nso_header.flags & NsoFlags_RoCompress) && out->nso_header.rodata_file_size != out->nso_header.rodata_segment_info.size) || \
        (out->nso_header.rodata_segment_info.file_offset + out->nso_header.rodata_file_size) > pfs_entry->size)
    {
        LOG_MSG("Invalid .rodata segment offset/size for NSO \"%s\"! (0x%08X, 0x%08X, 0x%08X).", out->nso_filename, out->nso_header.rodata_segment_info.file_offset, out->nso_header.rodata_file_size, \
                out->nso_header.rodata_segment_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.data_segment_info.file_offset < sizeof(NsoHeader) || !out->nso_header.data_segment_info.size || \
        ((out->nso_header.flags & NsoFlags_DataCompress) && (!out->nso_header.data_file_size || out->nso_header.data_file_size > out->nso_header.data_segment_info.size)) || \
        (!(out->nso_header.flags & NsoFlags_DataCompress) && out->nso_header.data_file_size != out->nso_header.data_segment_info.size) || \
        (out->nso_header.data_segment_info.file_offset + out->nso_header.data_file_size) > pfs_entry->size)
    {
        LOG_MSG("Invalid .data segment offset/size for NSO \"%s\"! (0x%08X, 0x%08X, 0x%08X).", out->nso_filename, out->nso_header.data_segment_info.file_offset, out->nso_header.data_file_size, \
                out->nso_header.data_segment_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.module_name_size > 1 && (out->nso_header.module_name_offset < sizeof(NsoHeader) || (out->nso_header.module_name_offset + out->nso_header.module_name_size) > pfs_entry->size))
    {
        LOG_MSG("Invalid module name offset/size for NSO \"%s\"! (0x%08X, 0x%08X).", out->nso_filename, out->nso_header.module_name_offset, out->nso_header.module_name_size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.api_info_section_info.size && (out->nso_header.api_info_section_info.offset + out->nso_header.api_info_section_info.size) > out->nso_header.rodata_segment_info.size)
    {
        LOG_MSG("Invalid .api_info section offset/size for NSO \"%s\"! (0x%08X, 0x%08X).", out->nso_filename, out->nso_header.api_info_section_info.offset, out->nso_header.api_info_section_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.dynstr_section_info.size && (out->nso_header.dynstr_section_info.offset + out->nso_header.dynstr_section_info.size) > out->nso_header.rodata_segment_info.size)
    {
        LOG_MSG("Invalid .dynstr section offset/size for NSO \"%s\"! (0x%08X, 0x%08X).", out->nso_filename, out->nso_header.dynstr_section_info.offset, out->nso_header.dynstr_section_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    if (out->nso_header.dynsym_section_info.size && (out->nso_header.dynsym_section_info.offset + out->nso_header.dynsym_section_info.size) > out->nso_header.rodata_segment_info.size)
    {
        LOG_MSG("Invalid .dynsym section offset/size for NSO \"%s\"! (0x%08X, 0x%08X).", out->nso_filename, out->nso_header.dynsym_section_info.offset, out->nso_header.dynsym_section_info.size);
        dump_nso_header = true;
        goto end;
    }
    
    /* Get module name. */
    if (!nsoGetModuleName(out)) goto end;
    
    /* Get .rodata segment. */
    if (!(rodata_buf = nsoGetRodataSegment(out))) goto end;
    
    /* Get module info name. */
    if (!nsoGetModuleInfoName(out, rodata_buf)) goto end;
    
    /* Get .api_info section data. */
    if (!nsoGetSectionFromRodataSegment(out, rodata_buf, (u8**)&(out->rodata_api_info_section), out->nso_header.api_info_section_info.offset, out->nso_header.api_info_section_info.size)) goto end;
    out->rodata_api_info_section_size = out->nso_header.api_info_section_info.size;
    
    /* Get .dynstr section data. */
    if (!nsoGetSectionFromRodataSegment(out, rodata_buf, (u8**)&(out->rodata_dynstr_section), out->nso_header.dynstr_section_info.offset, out->nso_header.dynstr_section_info.size)) goto end;
    out->rodata_dynstr_section_size = out->nso_header.dynstr_section_info.size;
    
    /* Get .dynsym section data. */
    if (!nsoGetSectionFromRodataSegment(out, rodata_buf, &(out->rodata_dynsym_section), out->nso_header.dynsym_section_info.offset, out->nso_header.dynsym_section_info.size)) goto end;
    out->rodata_dynsym_section_size = out->nso_header.dynsym_section_info.size;
    
    success = true;
    
end:
    if (rodata_buf) free(rodata_buf);
    
    if (!success)
    {
        if (dump_nso_header) LOG_DATA(&(out->nso_header), sizeof(NsoHeader), "NSO header dump:");
        
        nsoFreeContext(out);
    }
    
    return success;
}

static bool nsoGetModuleName(NsoContext *nso_ctx)
{
    if (nso_ctx->nso_header.module_name_offset < sizeof(NsoHeader) || nso_ctx->nso_header.module_name_size <= 1) return true;
    
    NsoModuleName module_name = {0};
    
    /* Get module name. */
    if (!pfsReadEntryData(nso_ctx->pfs_ctx, nso_ctx->pfs_entry, &module_name, sizeof(NsoModuleName), nso_ctx->nso_header.module_name_offset))
    {
        LOG_MSG("Failed to read NSO \"%s\" module name length!", nso_ctx->nso_filename);
        return false;
    }
    
    /* Verify module name length. */
    if (module_name.name_length != ((u8)nso_ctx->nso_header.module_name_size - 1))
    {
        LOG_MSG("NSO \"%s\" module name length mismatch! (0x%02X != 0x%02X).", nso_ctx->nso_filename, module_name.name_length, (u8)nso_ctx->nso_header.module_name_size - 1);
        return false;
    }
    
    /* Allocate memory for the module name. */
    nso_ctx->module_name = calloc(nso_ctx->nso_header.module_name_size, sizeof(char));
    if (!nso_ctx->module_name)
    {
        LOG_MSG("Failed to allocate memory for NSO \"%s\" module name!", nso_ctx->nso_filename);
        return false;
    }
    
    /* Read module name string. */
    if (!pfsReadEntryData(nso_ctx->pfs_ctx, nso_ctx->pfs_entry, nso_ctx->module_name, module_name.name_length, nso_ctx->nso_header.module_name_offset + 1))
    {
        LOG_MSG("Failed to read NSO \"%s\" module name string!", nso_ctx->nso_filename);
        return false;
    }
    
    return true;
}

static u8 *nsoGetRodataSegment(NsoContext *nso_ctx)
{
    int lz4_res = 0;
    bool compressed = (nso_ctx->nso_header.flags & NsoFlags_RoCompress), verify = (nso_ctx->nso_header.flags & NsoFlags_RoHash);
    
    u8 *rodata_buf = NULL;
    u64 rodata_buf_size = (compressed ? LZ4_DECOMPRESS_INPLACE_BUFFER_SIZE(nso_ctx->nso_header.rodata_segment_info.size) : nso_ctx->nso_header.rodata_segment_info.size);
    
    u8 *rodata_read_ptr = NULL;
    u64 rodata_read_size = (compressed ? nso_ctx->nso_header.rodata_file_size : nso_ctx->nso_header.rodata_segment_info.size);
    
    u8 rodata_hash[SHA256_HASH_SIZE] = {0};
    
    bool success = false;
    
    /* Allocate memory for the .rodata buffer. */
    if (!(rodata_buf = calloc(rodata_buf_size, sizeof(u8))))
    {
        LOG_MSG("Failed to allocate 0x%lX bytes for the .rodata segment in NSO \"%s\"!", rodata_buf_size, nso_ctx->nso_filename);
        return NULL;
    }
    
    rodata_read_ptr = (compressed ? (rodata_buf + (rodata_buf_size - nso_ctx->nso_header.rodata_file_size)) : rodata_buf);
    
    /* Read .rodata segment data. */
    if (!pfsReadEntryData(nso_ctx->pfs_ctx, nso_ctx->pfs_entry, rodata_read_ptr, rodata_read_size, nso_ctx->nso_header.rodata_segment_info.file_offset))
    {
        LOG_MSG("Failed to read .rodata segment in NRO \"%s\"!", nso_ctx->nso_filename);
        goto end;
    }
    
    if (compressed)
    {
        /* Decompress .rodata segment in-place. */
        if ((lz4_res = LZ4_decompress_safe((char*)rodata_read_ptr, (char*)rodata_buf, (int)nso_ctx->nso_header.rodata_file_size, (int)rodata_buf_size)) != \
            (int)nso_ctx->nso_header.rodata_segment_info.size)
        {
            LOG_MSG("LZ4 decompression failed for NRO \"%s\"! (0x%08X).", nso_ctx->nso_filename, (u32)lz4_res);
            goto end;
        }
    }
    
    if (verify)
    {
        /* Verify .rodata segment hash. */
        sha256CalculateHash(rodata_hash, rodata_buf, nso_ctx->nso_header.rodata_segment_info.size);
        if (memcmp(rodata_hash, nso_ctx->nso_header.rodata_segment_hash, SHA256_HASH_SIZE) != 0)
        {
            LOG_MSG(".rodata segment checksum mismatch for NRO \"%s\"!", nso_ctx->nso_filename);
            goto end;
        }
    }
    
    success = true;
    
end:
    if (!success && rodata_buf)
    {
        free(rodata_buf);
        rodata_buf = NULL;
    }
    
    return rodata_buf;
}

static bool nsoGetModuleInfoName(NsoContext *nso_ctx, u8 *rodata_buf)
{
    NsoModuleInfo *module_info = (NsoModuleInfo*)(rodata_buf + 0x4);
    if (!module_info->name_length) return true;
    
    /* Allocate memory for the module info name. */
    nso_ctx->module_info_name = calloc(module_info->name_length + 1, sizeof(char));
    if (!nso_ctx->module_info_name)
    {
        LOG_MSG("Failed to allocate memory for NSO \"%s\" module info name!", nso_ctx->nso_filename);
        return false;
    }
    
    /* Copy module info name. */
    sprintf(nso_ctx->module_info_name, "%.*s", (int)module_info->name_length, module_info->name);
    
    return true;
}

static bool nsoGetSectionFromRodataSegment(NsoContext *nso_ctx, u8 *rodata_buf, u8 **section_ptr, u64 section_offset, u64 section_size)
{
    if (!section_size || (section_offset + section_size) > nso_ctx->nso_header.rodata_segment_info.size) return true;
    
    /* Allocate memory for the desired .rodata section. */
    if (!(*section_ptr = malloc(section_size)))
    {
        LOG_MSG("Failed to allocate 0x%lX bytes for section at .rodata offset 0x%lX in NSO \"%s\"!", section_size, section_offset, nso_ctx->nso_filename);
        return false;
    }
    
    /* Copy .rodata section data. */
    memcpy(*section_ptr, rodata_buf + section_offset, section_size);
    
    return true;
}