IIPU.cpp

#include "IIPU.h"
const char* action = NULL;

static std::vector<uint8_t> Unpack(
    const std::string& sourceFileName,
    uint32_t dataOffset,
    uint32_t size
) {

    std::ifstream sourceFile(sourceFileName, std::ios::binary);
    if (!sourceFile.is_open()) {
        std::cerr << "Error opening input file: " << sourceFileName << std::endl; exit(1);
    }

    sourceFile.seekg(dataOffset);
    std::vector<uint8_t> sourceBuffer(size); // Это запакованные данные
    sourceFile.read((char*)sourceBuffer.data(), size);
    sourceFile.close();

    std::vector<uint8_t> outputBuffer(size); // Это распакованные данные

    uint32_t index = 0; // Текущее смещение в sourceFile
    bool exceptionTrigger = false;

    // Пропуск функции, если размер файла на входе равен 0
    if (size == 0) {
        std::cerr << "Input file is empty: " << sourceFileName << std::endl;
        outputBuffer.clear();
        return outputBuffer;
    }

    // Алгоритм распаковки
    while (true) {
        if (index >= size)
        {
            break;
        }

        if (sourceBuffer[index] == 0) {
            if (exceptionTrigger == false) {
                std::cerr << "INI_PACK is possibly corrupted, found null-byte at offset 0x" << std::hex << index << std::endl << "Continuing nevertheless..." << std::endl;
                exceptionTrigger = true;
            }
            index++;
            continue;
        }

        outputBuffer[index] = sourceBuffer[index] - sourceBuffer[index] / 2;

        index++;
    }
    action = "Unpack";

    return outputBuffer;
}

static std::vector<uint8_t> Compress(
    const std::string& sourceFileName,  //argv[4]
    uint32_t unpackedSize,
    uint32_t& packedSize
)
{
    std::ifstream sourceFile(sourceFileName, std::ios::binary | std::ios::ate); // Модификатор ate is crucial
    if (!sourceFile.is_open()) {
        std::cerr << "Error opening source file on compression: " << sourceFileName << std::endl; exit(1);
    }

    std::vector<uint8_t> buffer;
    std::streamsize fileSize = sourceFile.tellg();

    // Если файл пустой, превышает исходный размер или ошибка чтения размера
    if ((fileSize <= 0) || (fileSize > unpackedSize)) {
        buffer.clear();
        return buffer;
    }

    buffer.resize(static_cast<uint32_t>(unpackedSize)); // Выделение памяти для sourceBuffer равное fileSize
    sourceFile.seekg(0, std::ios::beg);
    sourceFile.read((char*)buffer.data(), unpackedSize); // Сохранение данных sourceFile в sourceBuffer
    sourceFile.close(); // Buffer получен, можно закрывать файл

    uint32_t index = 0; // Текущее смещение в sourceFile

    // Алгоритм запаковки
    while (true) {
        if (index >= fileSize)
        {
            break;
        }

        if (buffer[index] > 0x80) {
            std::cerr << "Invalid ini, found forbidden char at offset " << std::hex << index << std::endl;
            buffer.clear();
            return buffer;
        }

        buffer[index] = buffer[index] + buffer[index];

        index++;
    }

    action = "Compress";
    packedSize = fileSize; // Сжатие не применяется, поэтому размеры равны
    return buffer;
}

int main(int argc, char* argv[])
{
    system("cls");

    // Обработка аргументов, получение имения файла в fileName
    if (argc < 2 && 4 > argc) {
        printf("\rUsage: xxx.efi [action] <INI file>\n\nActions:\n\t-r\t\tReplace INI_PACK contents\n\n");
        return 1;
    }

    // Debug
    //printf("Args count: %d\n", argc);

    std::cout << "Loading..." << std::endl;

    // Общие переменные
    const std::string fileName = argv[1];
    action = argv[2];

    uint32_t align = NULL;
    uint32_t rawSize = NULL;
    std::vector<uint8_t> unpackedData;

    std::vector<uint32_t> offsets = { 0 };
    std::string signature = (std::string)"24494E495F5041434B"; // Сигнатура $INI_PACK
    uint8_t sigCount = IFLASHcountHexPattern(fileName, signature, offsets);

    // Переменные для Replace
    uint32_t packedSize = NULL;
    uint32_t packedDiff = NULL;
    std::vector<uint8_t> efiFileData;
    std::vector<uint8_t> packedData;

    if (sigCount == 0) {
        std::cerr << "There's something wrong with the input file... INI_PACK is missing" << std::endl;
        return 1;
    }
    
    
    std::ifstream sourceFile(fileName, std::ios::binary);

    if (!sourceFile.is_open()) {
        std::cerr << "Error opening source file while getting INI_PACK size: " << fileName << std::endl;
        return 1;
    }

    sourceFile.seekg(offsets[0] + 0xD); // 0xD - относительное смещение значения dataSize
    sourceFile.read(reinterpret_cast<char*>(&rawSize), 0x4); // Получить rawSize
    sourceFile.seekg(offsets[0] + 0x9); // 0x9 - относительное смещение значения align
    sourceFile.read(reinterpret_cast<char*>(&align), 0x4); // Получить align

    printf("\nSuccessfully found INI_PACK at offset: 0x%X\nData size: 0x%X\n", (int)offsets[0], rawSize);
    

    // Проверка сколько сигнатур нашло
    if (!(sigCount > 1)) {
        switch (argc) {
        case 2:

            // Сохранить запакованные данные
            saveDataByOffset(fileName, (std::string)"INI_PACK.bin", offsets[0] + align, rawSize);

            // Создать распакованный файл
            unpackedData = Unpack(fileName, offsets[0] + align, rawSize);
            saveBufferToFile((std::string)"INI_PACK.txt", unpackedData, rawSize);

            break;
        case 4:
            // Процедура запаковки и замены
            if (action != (std::string)"-r") { action = NULL; break; }

            packedData = Compress(argv[3], rawSize, packedSize);
            if (packedData.size() <= 1) { printf("Failed to compress the file: empty data\n"); return 1; }

            // Сохранение всего efi файла в буфер
            saveFileToBuffer(fileName, efiFileData);

            // Стереть старый модуль из буфера
            for (uint32_t i = 0; i < rawSize; i++) {
                efiFileData.erase(efiFileData.begin() + offsets[0] + align);
            }

            // Вставить запакованный файл
            for (size_t i = packedData.size() - 1; i > 0; i--) {
                efiFileData.insert(efiFileData.begin() + offsets[0] + align, packedData[i]);
            }
            efiFileData.insert(efiFileData.begin() + offsets[0] + align, packedData[0]);

            // Простое выявление большего значения и присвоение разницы
            packedDiff = rawSize >= packedData.size() ? (rawSize - packedData.size()) : NULL;

            // Debug
            //printf("packedDiff 0x%X - 0x%X\n", rawSize, packedData.size());

            // Убрать разницу в размерах. "(fileSize - romBlockSize)" это вычисление смещения romBlock. 4 - отступ для chksum bytes
            if (rawSize >= packedData.size()) {

                //Увеличить Rom, т.к. размер ориг ini, был больше чем размер данных для вставки (packedData.size)
                for (uint32_t i = 0; i < packedDiff; i++) {
                    efiFileData.insert(efiFileData.begin() + offsets[0] + align + packedData.size(), 0x0);
                }
            }
            else
            {
                printf("\nUh-oh, the new INI is so big it overpals with other data! Offset: 0x%X.\nNot doing replacement\n", offsets[0] + align + packedSize + 1);
                return 1;
            }

            std::cout << std::endl << "Saving re-packed INI to " << "new" << fileName << "..." << std::endl;
            saveBufferToFile((std::string)"new" + fileName, efiFileData, efiFileData.size());

            action = "Replace";
            break;
        default:
            // Если кол-во аргументов ни одно из заявленных
            std::cerr << std::endl << "Incorrect command!" << std::endl;
            return 1;
        }
        if (action == NULL) {
            std::cerr << std::endl << "Command failed!" << std::endl;
            return 1;
        }
    }
    else
    {
        std::cerr << "There's more than 1 INI_PACK in the source file" << std::endl << "Unable to pick one" << std::endl;
        return 1;
    }

    std::cout << std::endl << action <<" performed successfully!" << std::endl; //Если дошло до конца, вывод текста
    return 0;
}