/*!
 *
 * Copyright (c) 2023 Diality Inc. - All Rights Reserved.
 *
 * \copyright
 * THIS CODE MAY NOT BE COPIED OR REPRODUCED IN ANY FORM, IN PART OR IN
 * WHOLE, WITHOUT THE EXPLICIT PERMISSION OF THE COPYRIGHT OWNER.
 *
 * \file    Package.cpp
 * \author  (last)      Phil Braica
 * \date    (last)      23-Jan-2023
 * \author  (original)  Phil Braica
 * \date    (original)  23-Jan-2023
 */

#include "Package.h"
#include "Obfuscate.h"
#include "IDataProvider.h"
#include "IEnterBootLoader.h"
#include "UiSwUpdate.h"
#include "UiProtocol.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <thread>

// For Debug:
// #include <errno.h>
// #include <qdebug.h>

/*!
 * Must be greater than a signature, and big enough to be good for streaming.
 * 
 * A signature is base64'd and may be padded but it is 3/2 * key length in bytes plus the start/end text.
 * We expect 344 + 34 * 2 = 412 for the buffer.
 * 
 * For streaming it makes sense to use at least one memory block (4096)
 */
#define DEFAULT_BUF_SIZE    (8192)      

#define MIN_SIGNATURE_SIZE  (344)       ///< Min size of a signature, can be a little bigger (padding, whitespace).
#define MAX_WHITESPACE_SIZE (40)        ///< Max whitespace.

#define MAX_RECURSION   (3)     ///< We assume someone might put the file within a few folders on the USB by accident.

#define DELETE_RETRIES  (4)     ///< Upto 4 attempts at deleting a file if there is a race condition.

const std::string g_sig_start = "SIGNATURE_START===================";   ///< Start signature.
const std::string g_sig_end = "SIGNATURE_END===================";   ///< End signature.


/*!
 * \brief   Constructor.
 */
Package::Package() :
    _items(),
    _rsa(),
    _packageName(""),
    _packageVers(""),
    _signature(""),
    _xml(""),
    _info(""),
    _startXml(0),
    _startData(0),
    _pFileUpdating(nullptr),
    _pIEnterBL(nullptr) {

    ; // NOP.
}

/*!
 * \brief   Virtual destructor..
 */
Package::~Package() {
    this->abort();
}

/*!
 * \brief   Parse the file.
 * 
 * Note that key_location default to "", which means use the already loaded key.
 * This is useful for scanning files in multiple directories such as checking a USB file.
 * 
 * \param   fileName        Filename.
 * \param   key_location    The folder containing the public key, can contain that file's name.
 */
bool Package::parse(const std::string& fileName, const std::string& key_location) {
    
    bool verifiedOk = false;
    FILE* fp = NULL;
    _lastFileName = fileName;
    try 
    {
        clear();
        fp = fopen(fileName.c_str(), "rb");
        if (fp == NULL) {
            _info = "Could not open binary file: " + fileName + "\n";
            return false;
        }
        if (key_location.length() != 0) {
            _rsa.load_keys(key_location);
        }
        verifiedOk = parseInternal(fp);
    } catch(...) {
        ; // NOP.
    }
    
    if (fp != NULL) {
        fclose(fp);
    }
    return verifiedOk;
}

/*!
 * \brief   What are the targets available for the just parsed file? 
 *
 * \return  Vector of files that are parsable.
 */
std::vector<SwUpdateTargetEnum> Package::targetsAvailable() {
    std::vector<SwUpdateTargetEnum> rv;
    for (PackageItem pi : _items) {
        SwUpdateTargetEnum update_type =
            pi._fileType == "HD" ? HD :
            pi._fileType == "HDFPGA" ? HDFPGA :
            pi._fileType == "DG" ? DG :
            pi._fileType == "DGFPGA" ? DGFPGA :
            pi._fileType == "UI" ? UI : LINUX;

        if (std::find(rv.begin(), rv.end(), update_type) == rv.end()) {
            rv.push_back(update_type);
        }
    }
    return rv;
}

/*!
 * \brief   Start updating, key is a string.
 *
 * \param   fileName      File name to verify then use.
 * \param   public_key    The public key file content.
 * \param   update        The desired update targets, if empty use ALL.
 *
 * \return  True on successfully started.
 */
bool Package::start_hasKey(
    const std::string& fileName,
    const std::vector<SwUpdateTargetEnum>& update)
{
    _lastFileName = fileName;
    this->abort();  // Blocks till anything in progress is finished.
    clear();

    _pFileUpdating = NULL;    
    try {
        // Open the file, keep it open during the update process.
        // 
        // This prevents any other process or driver from altering the file 
        // during the update. It is a cyber security feature.
        _pFileUpdating = fopen(fileName.c_str(), "rb");
        if (_pFileUpdating == NULL) {
            int err = errno;
            _info = "Could not open binary file: " + fileName + ", errno=" + std::to_string(err) + "\n";
            return false;
        }

        // Now do a verify. This ensures that nothing altered the file between the time of
        // the user selecting it from the UI and the application and made it invalid.
        bool verifiedOk = parseInternal(_pFileUpdating);
        if (verifiedOk) {

            // Go through the items, create data providers.
            for (const PackageItem& pi : _items) {
                SwUpdateTargetEnum update_type =
                    pi._fileType == "HD" ? HD :
                    pi._fileType == "HDFPGA" ? HDFPGA :
                    pi._fileType == "DG" ? DG :
                    pi._fileType == "DGFPGA" ? DGFPGA :
                    pi._fileType == "UI" ? UI : LINUX;
                bool isFileCopy = (update_type == UI) || (update_type == LINUX);
                if ((std::find(update.begin(), update.end(), update_type) != update.end()) || (update.size() == 0)) {
                    SwUpdate::DataProvider_File dp(
                        _pFileUpdating,
                        update_type,
                        (uint32)(this->_startData + pi._byteOffset),
                        pi._size,
                        pi._security,
                        pi._version,
                        isFileCopy ? pi._destPath : "");

                    _streams.push_back(dp);
                }
            }
        }
    } catch(...) {
        if (_pFileUpdating != NULL) {
            fclose(_pFileUpdating);
            _pFileUpdating = NULL;
        }
        return false;
    }
    
    std::vector<SwUpdate::IDataProvider*> sv;
    for (std::size_t ii = 0; ii < _streams.size(); ii++) {
        sv.push_back(&_streams[ii]);
    }
    return SwUpdate::UiSwUpdate::instance().start(sv, _pIEnterBL);
}

/*!
 * \brief   Start updating.
 *
 * \param   fileName        File name to verify then use.
 * \param   key_location    The folder containing the public key, can contain that file's name.
 * \param   update          The desired update targets, if empty use ALL.
 *
 * \return  True on successfully started.
 */
bool Package::start(
    const std::string & fileName, 
    const std::string& key_location, 
    const std::vector<SwUpdateTargetEnum>& update) {

    _lastFileName = fileName;
    this->abort();  // Blocks till anything in progress is finished.
    clear();

    _pFileUpdating = fopen(fileName.c_str(), "rb");
    if (_pFileUpdating == NULL) {
        int err = errno;
        _info = "Could not open binary file: " + fileName + ", errno=" + std::to_string(err) + "\n";
        return false;
    }
    _rsa.load_keys(key_location);
    bool verifiedOk = parseInternal(_pFileUpdating);
    if (!verifiedOk) {
        fclose(_pFileUpdating);
        _pFileUpdating = NULL;
        _info += "Could not verify update package integrity: " + fileName + "\n";
        return false;
    }
    for (const PackageItem& pi : _items) {
        SwUpdateTargetEnum update_type =
            pi._fileType == "HD" ? HD :
            pi._fileType == "HDFPGA" ? HDFPGA :
            pi._fileType == "DG" ? DG :
            pi._fileType == "DGFPGA" ? DGFPGA :
            pi._fileType == "UI" ? UI : LINUX;
        bool isFileCopy = (update_type == UI) || (update_type == LINUX);
        if ((std::find(update.begin(), update.end(), update_type) != update.end()) || (update.size() == 0)) {
            SwUpdate::DataProvider_File dp(
                _pFileUpdating,
                update_type, 
                (uint32)(this->_startData + pi._byteOffset),
                pi._size,
                pi._security,
                pi._version,
                isFileCopy ? pi._destPath : "");

            _streams.push_back(dp);
        }
    }
    std::vector<SwUpdate::IDataProvider*> sv;
    for (std::size_t ii = 0; ii < _streams.size(); ii++) {
        sv.push_back(&_streams[ii]);
    }
    
    return SwUpdate::UiSwUpdate::instance().start(sv, _pIEnterBL);
}

/*!
 * \brief   Is it completed?
 *
 * \return  True if completed.
 */
bool Package::completedAll() {
    bool completed_all = SwUpdate::UiSwUpdate::instance().completedAll();
    if (completed_all) {
        _streams.clear();

        // Close file pointer.
        if (_pFileUpdating != nullptr) {
            try {
                fclose(_pFileUpdating);
            }
            catch (...) {
                ; // NOP.
            }
            _pFileUpdating = nullptr;
        }
    }
    return completed_all;
}

/*!
 * \brief   Abort update.
 */
void Package::abort() {
    SwUpdate::UiSwUpdate::instance().abort();

    // Now safe to discard these stream objects.
    _streams.clear();   

    // Close file pointer.
    if (_pFileUpdating != nullptr) {
        try {
            fclose(_pFileUpdating);
        }
        catch (...) {
            ; // NOP.
        }
        _pFileUpdating = nullptr;
    }
}

/*!
 * \brief   Current estimated progress.
 *
 * \return  Vector of statuses.
 */
std::vector<SwUpdate::UiUpdateStatus> Package::progress() {
    bool completed_all = SwUpdate::UiSwUpdate::instance().completedAll();

    if (completed_all) {
        _streams.clear();

        // Close file pointer.
        if (_pFileUpdating != nullptr) {
            try {
                fclose(_pFileUpdating);
            }
            catch (...) {
                ; // NOP.
            }
            _pFileUpdating = nullptr;
        }
    }

    return SwUpdate::UiSwUpdate::instance().progress();
}

/*!
 * \brief   Parse the file.
 *
 * \param   pFile           File object.
 */
bool Package::parseInternal(FILE* pFile) {

    // Update status.
    _info = "Reading ..\n";

    // Read the key.
    if (_rsa.public_exists() == false) {
        _info += "    Could not open key file.\n";
        return false;
    }
    else {
        _info += "    Read key file.\n";
    }

    // Get the signature.
    std::string signature = get_signature(pFile);
    if (signature.length() == 0) {
        return false;
    }

    // Verify the file.
    fseek(pFile, (long)_startXml, 0);
    bool verified_ok = _rsa.verify_file(pFile, signature);

    if (verified_ok) {
        updateCatalog(pFile);
    }
    else {

        this->_info += "    Verified failed.\n";
    }
    return verified_ok;
}

/*!
 * \brief   Clear our state out.
 */
void Package::clear() {
    this->_items.clear();
    this->_packageName = "";
    this->_packageVers = "";
    this->_signature = "";
    this->_xml = "";
    this->_startXml = 0;
    this->_startData = 0;
}

/*!
 * \brief   Get the signature.
 * 
 * \param   pFile   File pointer.
 * 
 * \return  Signature string.
 */
std::string Package::get_signature(FILE* pFile) {
    char buf[DEFAULT_BUF_SIZE];
    std::size_t sz = fread(buf, 1, DEFAULT_BUF_SIZE, pFile);
    std::string signature(buf, sz);
    std::size_t sig_start = signature.find(g_sig_start);

    if (sig_start == std::string::npos) {
        _info += "    Could not find signature start\n.";
        return "";
    }
    sig_start += g_sig_start.size();
    std::size_t sig_end = signature.find(g_sig_end, sig_start);
    if (sig_end == std::string::npos) {
        _info += "    Could not find signature end\n.";
        return "";
    }
    signature = signature.substr(sig_start, sig_end - sig_start);
    if (signature.length() < MIN_SIGNATURE_SIZE) {
        _info += "    Signature was too small. \n";
        return "";
    }
    if (signature.length() > MIN_SIGNATURE_SIZE + MAX_WHITESPACE_SIZE) {
        _info += "    Signature was too big. \n";
        return "";
    }
    _info += "    Signature found ok, was " + std::to_string(signature.length()) + " bytes.\n";
    _startXml = sig_end + g_sig_end.size();
    return signature;
}


/*!
 * \brief   Get the signature.
 *
 * \param   pFile   File pointer.
 *
 * \return  Signature string.
 */
void Package::updateCatalog(FILE *pFile) {
    Obfuscate ob;
    bool isFirst = true;
    // bool found = false;
    std::string xml_str;
    fseek(pFile, (long)_startXml, 0);
    char pBuf[DEFAULT_BUF_SIZE];
    std::string endTag = "</root>";
    while (true) {
        std::size_t sz = fread(pBuf, 1, DEFAULT_BUF_SIZE, pFile);
        if (sz == 0)
            break;
        xml_str += ob.decode(pBuf, sz, isFirst);
        isFirst = false;
        std::size_t xml_end = xml_str.find(endTag);
        if (xml_end != std::string::npos) {
            //found = true;
            xml_str = xml_str.substr(0, xml_end + endTag.size());
            break;
        }
        if (sz < DEFAULT_BUF_SIZE) {
            break;
        }
    }

    this->_packageName = PackageItem::unwrap(xml_str, "name");
    this->_packageVers = PackageItem::unwrap(xml_str, "version");
    this->_xml = xml_str;
    this->_items = PackageItem::fromXml(xml_str);
    this->_startData = this->_startXml + xml_str.length() + Obfuscate::headerSize();
    this->_info += "    Found " + std::to_string(this->_items.size()) + " items.\n";
    this->_info += "    Verified ok.\n";
}

/**
 * \brief   Get the list of available things in the directory.
 * 
 * \param   directory    Directory to scan (not recursive).
 * 
 * \return  vector of tupple: filename, fullNameAndPath, packagename, version
 */
std::vector<PackageInfo> Package::get_available(const std::string& directory)
{
    std::vector<PackageInfo> rv;
    this->get_availableInternal(directory, rv);
    return rv;
}

/**
 * \brief   Get the list of available things in many directories.
 *
 * \param   directories    Directories to scan (not recursive).
 *
 * \return  vector of tupple: filename, fullNameAndPath, packagename, version
 */
std::vector<PackageInfo> Package::get_available(const std::vector<std::string>& directories)
{
    std::vector<PackageInfo> rv;
    for (const std::string directory : directories) {
        this->get_availableInternal(directory, rv);       
    }

    return rv;
}

/**
 * \brief   Copy / cache it (if not already in there).
 *
 * \param   fullNamePath    The package being installed.
 * \param   cacheDirectory  Cache location.
 * \param   maxCacheSize    Maximum number of cached packages in cacheDirectory.
 *
 * \return  True on success.
 */
bool Package::copyUpdate(const std::string& fullNamePath, const std::string& cacheDirectory, uint32 maxCacheSize)
{
    // Ensure the cache directory exists, if not try to create it.
    // If we can't create it, then just return false.
    try 
    { 
        if (std::filesystem::is_directory(cacheDirectory) == false)
        {
            std::filesystem::create_directory(cacheDirectory);
        }
    }
    catch (...)
    {
        return false;
    }

    // Nothing to do?
    if (fullNamePath.find(cacheDirectory) == 0) {
        // Already there, nothing to do.
        return true;
    }

    // Find the name.
    std::string fileName = "";    
    for (const PackageInfo& pi : _packages)
    {
        // Try to find that entry.
        if (pi.fullPath == fullNamePath)
        {
            fileName = pi.fileName;
            break;
        }
    }
    if (fileName.size() == 0)
    {
        // Race condition, no longer exists, nothing to do.
        return false;
    }

    // If inCache > maxCacheSize, remove the oldest till enough room.
    // 
    // We allow upto 4 removals because in theory an old version of 
    // software may have allowed more or less than
    // the latest limit.
    for (std::size_t ii = 0; ii < 4; ii++) 
    { 
        uint inCache = 0;
        std::vector<PackageInfo>::iterator oldest = _packages.end();

        for (std::vector<PackageInfo>::iterator it = _packages.begin();
            it != _packages.end();
            it++)
        {
            // Sum up the entries.
            if (it->fullPath.find(cacheDirectory) == 0)
            {
                if (inCache == 0)
                {
                    oldest = it;
                } else {
                    if (oldest->lastWritten > oldest->lastWritten) 
                    {
                        oldest = it;
                    }
                }
                inCache++;
            }
        }
        
        if ((inCache >= maxCacheSize) && (oldest != _packages.end()))
        {
            // Remove the oldest, we already know it is not the same as the selected.
            bool deleted = false;
            for (std::size_t jj = 0; jj < DELETE_RETRIES; jj++) {
                try 
                {   
                    int deletedVal = std::remove(oldest->fullPath.c_str());
                    if (deletedVal == 0) {
                        deleted = true;
                        break;
                    }

                    // Sleep 250 ms hoping what ever had this open releases it.
                    using namespace std::chrono_literals;
                    std::this_thread::sleep_for(250ms);
                } catch (...) {
                    // Fail as we might not have room for it!
                    return false;
                }
            }
            if (deleted) {
                _packages.erase(oldest);
            }
            else
            {
                break;
            }
        } else {
            break;
        }
    }

    // Now copy it.
    try
    {
        // There could be another file of the same name there, so we use a number to get a unique one.
        // 
        // As there can't be more than maxCacheSize there (assuming delete worked) we scan through 
        // no more than maxCacheSize*2.
        std::size_t dot = fileName.find('.');
        std::string baseName = fileName;
        std::string suffix = "";
        if (dot > 1) {
            baseName = fileName.substr(0, dot);
            suffix = fileName.substr(dot);
        }
        std::string newName = cacheDirectory + "/" + fileName;
        
        for (std::size_t ii = 0; ii < maxCacheSize * 2; ii++) {
            if (std::filesystem::exists(newName))
            {
                newName = cacheDirectory + "/" + baseName + "_" + std::to_string(ii) + suffix;
            } else {
                break;
            }
        }
        std::filesystem::copy_file(fullNamePath, newName);
    }
    catch (...) {
        // Fail as we might not have room for it!
        return false;
    }
    
    // Update _packages.
    for (PackageInfo& pi : _packages)
    {
        // Try to find that entry.
        if (pi.fullPath == fullNamePath)
        {
            pi.fullPath = cacheDirectory + "/" + fileName;
            pi.lastWritten = std::filesystem::last_write_time(pi.fullPath);
            break;
        }
    }
    return true;
}

/*!
 * \brief   If a script exists, there can be but one, get it's destination.
 * 
 * \return  Destination location of the script.
 */
std::string Package::getScriptIfAny() {
    for (const PackageItem& pi : _items) 
    {
        if (pi._isScript) {
            return pi._destPath;
        }
    }
    return "";
}

/*!
 * \brief   Reboot to run the firmware (both HD and DG) apps if they are valid.
 */
void Package::rebootFW() {
    //const int retryRebootEffort = 4;
    SwUpdate::MsgLink link(KBPS_WIRE);

    for (std::size_t ii = 0; ii < 2; ii++) {
        SwUpdateCommand msg;
        msg.id = SwUpdate::UiProtocol::getNextMsgSlotId();
        msg.cmd = SwUpdate_FormCommand(ii == 0 ? HD : DG, RunApp);
        msg.rand = SwUpdate::UiProtocol::getNextMsgSeed();
        SwUpdate_createSecurity((uint8 *)&msg, sizeof(SwUpdateCommand));
        // Doesn't build, inlined link call and it has a memcpy line 175 failure.
        //   link.sendOk(SwUpdateCanMsgIds::CommandId, (uint8 *)&msg, retryRebootEffort);
    }
}

/*!
 * \brief   Scan a directory recursively for regular files and avoid links.
 *
 * \param   directory       The directory to scan.
 * \param   depth           Depth remaining to scan.
 * \param   values          Structure to append into.
 *
 * \return  List of files (full path) that are regular files.
 */
void Package::recurseDirectory(
    const std::string& directory,
    int depth,
    std::vector<std::filesystem::directory_entry>& values)
{
    std::vector<std::string> rv;

    if (depth < 0) {
        return;
    }

    depth--;

    // Use try to ignore things at the directory and file level.
    // 
    // If one of the directories doesn't exist OR no permission,
    // then that's odd but we could be trying to restore via update
    // a semi-incompatible old/new version that has an odd directory
    // layout. By ignoring anything thrown, we let the user 
    // recover cleanly from that situation.
    try
    {
        if (std::filesystem::is_directory(directory)) {
            for (const std::filesystem::directory_entry& entry :
                std::filesystem::directory_iterator(directory)) {

                try
                {
                    // Ignore symlinks, cyber threat avoidance.
                    if (entry.is_symlink()) {
                        continue;
                    }

                    // Could be an update package, add it.
                    if (entry.is_regular_file()) {
                        values.push_back(entry);
                    }

                    // If allowed, follow it.
                    if (entry.is_directory() && (depth > 0)) {
                        recurseDirectory(entry.path().string(), depth, values);
                    }
                }
                catch (...) {
                    ; // Allow the user to continue.
                }
            }
        }
    }
    catch (...) {
        ; // Allow the user to continue.
    }
}

/*!
 * \brief   Scan a vector for a file by full name and path.
 * 
 * \param   packages        Vector to search.
 * \param   fullNamePath    Full name and path.
 * 
 * \return  > 0 if found, else -1.
 */
int Package::findPackageVector(
    const std::vector<PackageInfo>& packages, 
    const std::string& fullNamePath) 
{
    for (int ii = 0; ii < (int)packages.size(); ii++) {
        if (packages[ii].fullPath == fullNamePath)
        {
            return ii;
        }
    }
    return -1;
}

/**
 * \brief   Get the list of available things in a directory
 *
 * \param   directories     Directories to scan (not recursive).
 * \param   info            Vector to append to.
 */
void Package::get_availableInternal(const std::string& directory, std::vector<PackageInfo>& info)
{
    if (_rsa.public_exists()) {

        std::vector< std::filesystem::directory_entry> entries;
        this->recurseDirectory(directory, MAX_RECURSION, entries);

        for (const std::filesystem::directory_entry& entry : entries) {
            // Public key exists, worth trying ... 

            // Entry is a normal file, validate / check it.    
            std::string fileNamePath = entry.path().string();
            for (std::size_t ii = 0; ii < fileNamePath.size(); ii++) {
                fileNamePath[ii] = fileNamePath[ii] == 0x5c ? 0x2f : fileNamePath[ii];
            }

            // In our cache?
            int found = findPackageVector(_packages, fileNamePath);
            if (found >= 0) 
            {
                info.push_back(_packages[found]);
                continue;
            }

            if (std::find(_avoid.begin(), _avoid.end(), fileNamePath) != _avoid.end()) {
                // in the don't care list, 
                continue;
            }

            FILE* fp = NULL;
            try
            {
                fp = fopen(fileNamePath.c_str(), "rb");
                if (fp != NULL) {
                    bool ok = parseInternal(fp);
                    if (ok) {

                        std::filesystem::file_time_type ftime = std::filesystem::last_write_time(fileNamePath);

                        // It's validated...
                        // Get the filename to show to the user, the path, 
                        // the package name and the internal version info.
                        //
                        // known not to be in _packages so add it to both.
                        _packages.push_back({
                            fileNamePath,
                            entry.path().filename().string(),
                            _packageName,
                            _packageVers,
                            ftime,
                            ""});
                        
                        info.push_back({
                            fileNamePath,
                            entry.path().filename().string(),
                            _packageName,
                            _packageVers,
                            ftime,
                            ""});
                    }
                    else 
                    {
                        if (std::find(_avoid.begin(), _avoid.end(), fileNamePath) == _avoid.end()) {
                            _avoid.push_back(fileNamePath);
                            // qDebug() << "Not added " << QString::fromStdString(fileNamePath) << "\n";
                        }
                    }
                }
                else 
                {
                    // int err = errno;
                    // qDebug() << "Not opened " << QString::fromStdString(fileNamePath) << " errno = " << err << "\n";
                }
            } catch (...) {
                // qDebug() << "Throw 1\n";
                ; // Allow the user to keep iterating across the rest of this directory.
            }
            
            // Close it.
            if (fp != NULL) {
                int fclose_rv = fclose(fp);
                if (fclose_rv != 0) {
                    // int err = errno;
                    // qDebug() << "Not closed " << QString::fromStdString(fileNamePath) << " errno = " << err << " fcloserv = " << fclose_rv << "\n";
                }
            }
        }
    }
}