
/*
 * Copyright (C) Jan 2006 Mellanox Technologies Ltd. All rights reserved.
 * Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * This software is available to you under a choice of one of two
 * licenses.  You may choose to be licensed under the terms of the GNU
 * General Public License (GPL) Version 2, available from the file
 * COPYING in the main directory of this source tree, or the
 * OpenIB.org BSD license below:
 *
 *     Redistribution and use in source and binary forms, with or
 *     without modification, are permitted provided that the following
 *     conditions are met:
 *
 *      - Redistributions of source code must retain the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer.
 *
 *      - Redistributions in binary form must reproduce the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer in the documentation and/or other materials
 *        provided with the distribution.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
 * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 */

#include <dirent.h>
#include <stdio.h>
#include <stdlib.h>

#include "image_access.h"
#include "img_version.h"
#include <string.h>
#include <iostream>
#include <mfa.h>
#include <mlxfwops.h>
#include <pldm_pkg.h>
#include <pldm_dev_id_record.h>
#include <pldm_component_image.h>
#include <tools_layouts/image_info_layouts.h>

using namespace std;

ImageAccess::ImageAccess(int compareFFV)
{
    _imgFwOps = NULL;
    _compareFFV = compareFFV;
    memset(&_imgFwParams, 0, sizeof(_imgFwParams));
    memset(_errBuff, 0, sizeof(_errBuff));
}

ImageAccess::~ImageAccess()
{
    if (_imgFwOps)
    {
        _imgFwOps->FwCleanUp();
        delete _imgFwOps;
        _imgFwOps = NULL;
    }
}

bool ImageAccess::hasMFAs(string directory)
{
    DIR* d;
    struct dirent* dir;
    int rc;
    d = opendir(directory.c_str());
    if (d == NULL)
    {
        return false;
    }
    while ((dir = readdir(d)) != NULL)
    {
        string fl = dir->d_name;
        if (fl == "." || fl == "..")
        {
            continue;
        }
        string fpath = directory;
        fpath += "/";
        fpath += fl;
        rc = getFileSignature(fpath);
        if (rc == IMG_SIG_TYPE_BIN || rc == IMG_SIG_TYPE_MFA)
        {
            closedir(d);
            return true;
        }
    }
    closedir(d);
    return false;
}
int ImageAccess::queryDirPsid(string& path,
                              string& psid,
                              string& selector_tag,
                              int image_type,
                              vector<PsidQueryItem>& riv)
{
    int res = 0;
    DIR* d;
    struct dirent* dir;
    int found = 0;
    int rc;

    d = opendir(path.c_str());
    if (d == NULL)
    {
        return -1;
    }

    while ((dir = readdir(d)) != NULL)
    {
        PsidQueryItem ro;
        string fl = dir->d_name;
        if (fl == "." || fl == "..")
        {
            continue;
        }
        string fpath = path;
        fpath += "/";
        fpath += fl;

        rc = this->queryPsid(fpath.c_str(), psid, selector_tag, image_type, ro);
        if (rc < 0)
        {
            res = -1;
            goto clean_up;
        }
        if (rc == 1)
        {
            riv.push_back(ro);
            found++;
            riv[0].found = found;
            res = 1;
        }
    }

    if (found > 1)
    {
        res = -2;
    }

clean_up:
    closedir(d);
    return res;
}

bool ImageAccess::openImg(fw_hndl_type_t hndlType, char* psid, char* fileHndl)
{
    memset(_errBuff, 0, sizeof(_errBuff));
    _imgFwParams.errBuff = _errBuff;
    _imgFwParams.errBuffSize = MLNX_ERR_BUFF_SIZE;
    _imgFwParams.hndlType = hndlType;
    _imgFwParams.psid = psid;
    _imgFwParams.fileHndl = fileHndl;
    _imgFwParams.shortErrors = true;
    _imgFwOps = FwOperations::FwOperationsCreate(_imgFwParams);
    if (_imgFwOps == NULL)
    {
        _errMsg = "-E- " + (string)_errBuff + "\n";
        if (_errMsg.find("upgrade tool") != std::string::npos)
        {
            _warning = _errMsg;
        }
        else if (_errMsg.find("with higher Version") != std::string::npos)
        {
            _warning = _errMsg;
        }
        _log += _errMsg;
        return false;
    }
    return true;
}

int ImageAccess::queryPsid(const string& fname,
                           const string& psid,
                           string& selector_tag,
                           int image_type,
                           PsidQueryItem& ri)
{
    int res = 0;
    mfa_desc* mfa_d = NULL;
    fw_info_t img_query;
    string iniName = "";
    vector<u_int8_t> sect; // to get fw configuration.
    vector<u_int8_t> dest;

    if (!openImg(FHT_FW_FILE, (char*)psid.c_str(), (char*)fname.c_str()))
    {
        return 0;
    }

    ri.psid = psid;
    ri.found = 0;

    memset(&img_query, 0, sizeof(img_query));
    if (!_imgFwOps->FwQuery(&img_query, true))
    {
        _errMsg = "Failed to query " + (string)_imgFwOps->err();
        _log += _errMsg;
        goto clean_up;
    }

    if (!psid.compare(img_query.fw_info.psid))
    {
        u_int32_t* supporteHwId;
        u_int32_t supporteHwIdNum;
        _imgFwOps->getSupporteHwId(&supporteHwId, supporteHwIdNum);
        ri.found = 1;
        ri.url = fname;
        ri.pns = "";
        ri.devId = (supporteHwIdNum > 0) ? supporteHwId[0] : 0;
        ri.revId = (ri.devId >> 16) & 0xffff;
        ri.devId = ri.devId & 0xffff;
        ri.isFailSafe = img_query.fw_info.is_failsafe;
        if (_imgFwOps->FwGetSection(H_FW_CONF, sect))
        {
            if (unzipDataFile(sect, dest, "Fw Configuration"))
            {
                char* ptr = strstr((char*)&dest[0], "Name =");
                int counter = 7;
                while (ptr && *ptr != '\n' && *ptr != '\r')
                {
                    if (counter-- > 0)
                    {
                        ptr++;
                        continue;
                    }
                    iniName += *ptr;
                    ptr++;
                }
            }
        }
        ri.iniName = iniName;
        if (!mfa_open_file(&mfa_d, (char*)fname.c_str()))
        {
            char* pn = mfa_get_board_metadata(mfa_d, (char*)psid.c_str(), (char*)"PN");
            if (pn != NULL)
            {
                ri.pns = pn;
                free(pn);
            }
            char* desc =
              mfa_get_board_metadata(mfa_d, (char*)psid.c_str(), (char*)"DESCRIPTION"); // TODO: Make more efficient
            if (desc != NULL)
            {
                ri.description = desc;
                free(desc);
            }
            mfa_close(mfa_d);
        }
        if (!_compareFFV)
        {
            ImgVersion imgv;
            imgv.setVersion("FW", 3, img_query.fw_info.fw_ver, img_query.fw_info.branch_ver);
            ri.imgVers.push_back(imgv);
        }
        else
        {
            u_int16_t fwVer[4];
            fwVer[0] = img_query.fw_info.fw_ver[0];
            fwVer[1] = img_query.fw_info.fw_ver[1];
            fwVer[2] = img_query.fw_info.fw_ver[2] / 100;
            fwVer[3] = img_query.fw_info.fw_ver[2] % 100;
            ImgVersion imgv;
            imgv.setVersion("FW", 4, fwVer);
            ri.imgVers.push_back(imgv);
        }
        for (int i = 0; i < img_query.fw_info.roms_info.num_of_exp_rom; i++)
        {
            ImgVersion imgVer;
            const char* tpc = _imgFwOps->expRomType2Str(img_query.fw_info.roms_info.rom_info[i].exp_rom_product_id);
            if (tpc == NULL)
            {
                // imgVer.setExpansionRomtoUnknown();
                tpc = "UNKNOWN_ROM";
            }
            int sz = img_query.fw_info.roms_info.rom_info[i].exp_rom_num_ver_fields;
            imgVer.setVersion(tpc, sz, img_query.fw_info.roms_info.rom_info[i].exp_rom_ver);
            ri.imgVers.push_back(imgVer);
        }
        res = 1;
    }

clean_up:
    (void)selector_tag;
    (void)image_type;
    _imgFwOps->FwCleanUp();
    delete _imgFwOps;
    _imgFwOps = NULL;

    return res;
}

int ImageAccess::getImage(const string& fname,
                          const string& psid,
                          string& selector_tag,
                          int image_type,
                          u_int8_t** filebuf)
{
    int res = -1;
    int type = getFileSignature(fname);

    if (type == IMG_SIG_TYPE_MFA)
    {
        mfa_desc* mfa_d;
        int rc = mfa_open_file(&mfa_d, (char*)fname.c_str());
        if (!rc)
        {
            res = mfa_get_image(mfa_d, (char*)psid.c_str(), image_type, (char*)selector_tag.c_str(), filebuf);
            mfa_close(mfa_d);
        }
        else
        {
            res = -1;
        }
    }
    else if (type == IMG_SIG_TYPE_BIN)
    {
        res = getImage(fname, filebuf);
    }
    // else type is unknown and res = -1

    return res;
}

int ImageAccess::getImage(const string& fname, u_int8_t** filebuf)
{
    int res = -1;
    FILE* fp;
    u_int8_t* buf = NULL;
    long int fsize;
    long int sz;

    if ((fp = fopen(fname.c_str(), "rb")) == NULL)
    {
        _errMsg = "Failed to open file : " + (string)fname.c_str();
        _log += _errMsg;
        return -1;
    }

    if (fseek(fp, 0L, SEEK_END) < 0)
    {
        _errMsg = "Fseek failed on file : " + (string)fname.c_str();
        _log += _errMsg;
        goto clean_up;
    }

    fsize = ftell(fp);
    if (fsize <= 0)
    {
        _errMsg = "Ftell failed on file : " + (string)fname.c_str();
        _log += _errMsg;
        goto clean_up;
    }
    rewind(fp);

    buf = (u_int8_t*)malloc(fsize * sizeof(u_int8_t));
    if (buf == NULL)
    {
        _errMsg = "Memory allocation failure for file : " + (string)fname.c_str();
        _log += _errMsg;
        goto clean_up;
    }

    memset(buf, 0, fsize);
    *filebuf = buf;
    sz = fread(buf, 1, fsize, fp);
    if (sz != fsize || sz <= 0 || buf == NULL)
    {
        _errMsg = "ftell size is not equivalent to fread total size" + (string)fname.c_str();
        _log += _errMsg;
        goto clean_up;
    }

    res = fsize;

clean_up:
    fclose(fp);
    if (res == -1)
    {
        if (buf != NULL)
        {
            free(buf);
        }
    }
    return res;
}

int ImageAccess::getFileSignature(const string& fname)
{
    static const int header_size = 16;
    FILE* fin;
    char tmpb[header_size];
    int res = -2;
    do
    {
        if (!(fin = fopen(fname.c_str(), "r")))
        {
            break;
        }

        if (1 > fread(tmpb, header_size, 1, fin))
        {
            break;
        }
        res = getBufferSignature((u_int8_t*)tmpb, 4);
    } while (0);

    if (fin)
    {
        fclose(fin);
    }
    return res;
}

int ImageAccess::getBufferSignature(u_int8_t* buf, u_int32_t size)
{
    int res = IMG_SIG_TYPE_UNKNOWN;

    if (0 == memcmp(buf, "MTFW", 4))
    {
        res = IMG_SIG_TYPE_BIN;
    }
    else if (0 == memcmp(buf, "MFAR", 4))
    {
        res = IMG_SIG_TYPE_MFA;
    }
    else if (0 == memcmp(buf, PldmPkg::UUID, size))
    {
        res = IMG_SIG_TYPE_PLDM;
    }

    return res;
}

int ImageAccess::get_bin_content(const string& fname, vector<PsidQueryItem>& riv)
{
    fw_info_t img_query;
    int res = 0;
    char psid[33];
    PsidQueryItem item;
    ImgVersion imgv;

    if (_imgFwOps != NULL)
    {
        _errMsg = "_imgFwOPs should be null";
        _log += _errMsg;
        return -1;
    }
    if (!openImg(FHT_FW_FILE, NULL, (char*)fname.c_str()))
    {
        return -1;
    }

    memset(&img_query, 0, sizeof(img_query));
    if (!_imgFwOps->FwQuery(&img_query, true))
    {
        _errMsg = "Failed to query " + (string)_imgFwOps->err();
        _log += _errMsg;
        res = -1;
        goto clean_up;
    }

    strcpy(psid, img_query.fw_info.psid);
    psid[32] = 0;
    item.psid = psid;
    item.pns = "";
    item.board_rev = "";

    imgv.setVersion("FW", 3, img_query.fw_info.fw_ver, img_query.fw_info.branch_ver);
    item.imgVers.push_back(imgv);
    for (int i = 0; i < img_query.fw_info.roms_info.num_of_exp_rom; i++)
    {
        ImgVersion imgVer;
        const char* tpc = _imgFwOps->expRomType2Str(img_query.fw_info.roms_info.rom_info[i].exp_rom_product_id);
        if (tpc == NULL)
        {
            // imgVer.setExpansionRomtoUnknown();
            tpc = "UNKNOWN_ROM";
        }
        int sz = img_query.fw_info.roms_info.rom_info[i].exp_rom_num_ver_fields;
        if (img_query.fw_info.roms_info.rom_info[i].exp_rom_product_id == 0xf)
        {
            sz = 1;
        }
        imgVer.setVersion(tpc, sz, img_query.fw_info.roms_info.rom_info[i].exp_rom_ver);
        item.imgVers.push_back(imgVer);
    }
    riv.push_back(item);

clean_up:
    _imgFwOps->FwCleanUp();
    delete _imgFwOps;
    _imgFwOps = NULL;

    return res;
}

int ImageAccess::get_mfa_content(const string& fname, vector<PsidQueryItem>& riv)
{
    mfa_desc* mfa_d;
    map_entry_hdr* me = NULL;
    char psid[33];

    if (mfa_open_file(&mfa_d, (char*)fname.c_str()))
    {
        return -1;
    }
    me = NULL;
    while ((me = mfa_get_next_mentry(mfa_d, me)))
    {
        PsidQueryItem item;
        strncpy(psid, me->board_type_id, 32);
        psid[32] = 0;
        item.psid = psid;
        char* val = mfa_get_map_entry_metadata(me, (char*)"PN");
        if (val == NULL)
        {
            item.pns = "";
        }
        else
        {
            item.pns = val;
            free(val);
        }
        val = mfa_get_map_entry_metadata(me, (char*)"BOARD_REV");
        if (val == NULL)
        {
            item.board_rev = "";
        }
        else
        {
            item.board_rev = val;
            free(val);
        }
        val = mfa_get_map_entry_metadata(me, (char*)"DESCRIPTION");
        if (val == NULL)
        {
            item.description = "";
        }
        else
        {
            item.description = val;
            free(val);
        }
        val = mfa_get_map_entry_metadata(me, (char*)"BRANCH");
        if (val == NULL)
        {
            item.branch = "";
        }
        else
        {
            item.branch = val;
            free(val);
        }

        map_image_entry* imge;
        for (int i = 0; i < me->nimages; i++)
        {
            imge = mfa_get_map_image(me, i);
            if (imge == NULL)
            {
                continue;
            }
            toc_entry* toc = mfa_get_image_toc(mfa_d, imge);
            if (toc == NULL)
            {
                continue;
            }
            if (toc->num_ver_fields != 0)
            {
                item.type = imge->image_type;
                char select_tag[33];
                strncpy(select_tag, imge->select_tag, 32);
                select_tag[32] = 0;
                item.selector_tag = select_tag;
                char* subImgType;
                if ((subImgType = mfasec_get_sub_image_type_str(toc->subimage_type)) != NULL)
                {
                    ImgVersion imgv;
                    if (subImgType[0] != '\0')
                    { // Not a padding section
                        imgv.setVersion(subImgType, toc->num_ver_fields, toc->version, item.branch);
                        item.imgVers.push_back(imgv);
                    }
                }
            }
        }
        riv.push_back(item);
    }
    mfa_close(mfa_d);
    return 0;
}
#define ITOC_ASCII 0x49544f43
#define TOC_HEADER_SIZE 0x20
#define TOC_ENTRY_SIZE 0x20
#define FS3_DEFAULT_SECTOR_SIZE 0x1000

bool ImageAccess::extract_pldm_image_info(const u_int8_t* buff, u_int32_t size, PsidQueryItem& query_item)
{
    FImage fimage;
    if (!fimage.open((u_int32_t*)buff, size))
    {
        return false;
    }
    static const u_int32_t sector_size = FS3_DEFAULT_SECTOR_SIZE;
    u_int32_t offset = 0;
    offset = (offset % sector_size == 0) ? offset : (offset + sector_size - offset % 0x1000);
    u_int8_t buffer[TOC_HEADER_SIZE] = {0}, entry_buffer[TOC_ENTRY_SIZE] = {0};
    struct connectx4_itoc_header itoc_header;
    while (offset < fimage.get_size())
    {
        // read ITOC header
        fimage.read(offset, buffer, TOC_HEADER_SIZE);
        connectx4_itoc_header_unpack(&itoc_header, buffer);
        if (itoc_header.signature0 != ITOC_ASCII)
        {
            offset += sector_size;
            continue;
        }
        int section_index = 0;
        struct connectx4_itoc_entry toc_entry;
        do
        {
            // read ITOC entry
            u_int32_t entry_addr = offset + TOC_HEADER_SIZE + section_index * TOC_ENTRY_SIZE;
            fimage.read(entry_addr, entry_buffer, TOC_ENTRY_SIZE);

            connectx4_itoc_entry_unpack(&toc_entry, entry_buffer);
            if (toc_entry.type == FS3_IMAGE_INFO)
            {
                u_int32_t flash_addr = toc_entry.flash_addr << 2;
                u_int32_t entry_size_in_bytes = toc_entry.size * 4;
                u_int8_t* buff = new u_int8_t[entry_size_in_bytes];
                memset(buff, 0, entry_size_in_bytes);
                fimage.read(flash_addr, buff, entry_size_in_bytes);
                // read image info
                struct connectx4_image_info image_info;
                connectx4_image_info_unpack(&image_info, buff);
                query_item.name = image_info.name;
                query_item.pns = query_item.name;
                query_item.description = image_info.description;
                static const int FW_VER_SIZE = 3;
                u_int16_t fw_ver[FW_VER_SIZE];
                char fw_branch[BRANCH_LEN + 1] = {0};
                fw_ver[0] = image_info.FW_VERSION.MAJOR;
                fw_ver[1] = image_info.FW_VERSION.MINOR;
                fw_ver[2] = image_info.FW_VERSION.SUBMINOR;
                (strncpy(fw_branch, image_info.vsd, BRANCH_LEN));
                fw_branch[BRANCH_LEN] = (char)0;
                ImgVersion imgv;
                imgv.setVersion("FW", FW_VER_SIZE, fw_ver, fw_branch);
                query_item.imgVers.push_back(imgv);

                delete[] buff;
                break;
            }

            section_index++;
        } while (toc_entry.type != FS3_END);
        break;
    }
    fimage.close();
    return 0;
}

int ImageAccess::get_pldm_content(const string& fname, vector<PsidQueryItem>& riv)
{
    PldmBuffer pldm_buff;
    pldm_buff.loadFile(fname);
    PldmPkg pldm;
    pldm.unpack(pldm_buff);

    u_int8_t dev_count = pldm.getDeviceIDRecordCount();
    for (u_int8_t i = 0; i < dev_count; i++)
    {
        PldmDevIdRecord* rec = pldm.getDeviceIDRecord(i);
        PsidQueryItem item;
        item.psid = rec->getDevicePsid();
        item.description = rec->getDescription();
        int image_index = rec->getComponentImageIndex();
        PldmComponenetImage* image_obj = pldm.getComponentImage(image_index);
        extract_pldm_image_info(image_obj->getComponentData(), image_obj->getComponentSize(), item);
        riv.push_back(item);
    }

    return 0;
}

int ImageAccess::get_file_content(const string& fname, vector<PsidQueryItem>& riv)
{
    int type = getFileSignature(fname);
    int res = -1;
    switch (type)
    {
        case IMG_SIG_TYPE_BIN:
            res = get_bin_content(fname, riv);
            break;
        case IMG_SIG_TYPE_MFA:
            res = get_mfa_content(fname, riv);
            break;
        case IMG_SIG_TYPE_PLDM:
            res = get_pldm_content(fname, riv);
            break;
    }
    return res;
}

string ImageAccess::getLastErrMsg()
{
    return _errMsg;
}

string ImageAccess::getlastWarning()
{
    return _warning;
}

string ImageAccess::getLog()
{
    return _log;
}
