boothifier/webSock/AppUpgrade.cpp

#include "AppUpgrade.h"
#include "esp_log.h"
#include <MD5Builder.h>
#include <LittleFS.h>
#include <memory>
#include "global.h"
#include "jsonConstrain.h"

static const char* TAG = "AppUpdater";
TaskHandle_t Update_Task_Handle = NULL;

// Queue handle for firmware update messages
//QueueHandle_t updateMsgQueue = NULL;


bool Version::operator<(const Version& other) const {
    if (major != other.major) return major < other.major;
    if (minor != other.minor) return minor < other.minor;
    return patch < other.patch;
}

String Version::toString() const {
    return String(major) + "." + String(minor) + "." + String(patch);
}

AppUpdater::AppUpdater(fs::FS& fs, const char* currVersion, const char* bucket, const char* manifestName, const char* appBin)
    : currentVersion(currVersion), bucketUrl(bucket), manifestName(manifestName), appName(appBin), fileSystem(fs), downloadBuffer(new uint8_t[BUFFER_SIZE])
{
    ESP_LOGI(TAG, "AppUpdater initialized with version %s", currVersion);
}

void AppUpdater::setProgressCallback(void (*callback)( UpdateStatus status, int percentage, const char* message)) {
    progressCb = callback;
}

void AppUpdater::updateProgress(UpdateStatus newStatus, int percentage, const char* message) {
    status = newStatus;
    if (progressCb) {
        progressCb(status, percentage, message);
    }
}

bool AppUpdater::checkManifest() {
    String url = String(bucketUrl) + manifestName;
    ESP_LOGD(TAG, "Fetching manifest from: %s", url.c_str());

    // Start the HTTP client and Send GET request for manifest
    HTTPClient http;
    http.begin(url);
    int httpCode = http.GET();
    if (httpCode != HTTP_CODE_OK) {
        ESP_LOGE(TAG, "HTTP GET failed, error: %d", httpCode);
        http.end();
        return false;
    }

    // Read the response
    String payload = http.getString();
    http.end();

    // Parse JSON
    DeserializationError error = deserializeJson(jsonManifest, payload);
    if (error) {
        ESP_LOGE(TAG, "Failed to parse manifest: %s", error.c_str());
        return false;
    }

    // Check for version section
    JsonObject jsonVersion = jsonManifest["version"];
    if (jsonVersion.isNull()) {
        ESP_LOGE(TAG, "No version section in manifest");
        return false;
    }

    // Get the remote version
    int major = jsonVersion["major"] | 0;
    int minor = jsonVersion["minor"] | 0;
    int patch = jsonVersion["patch"] | 0;
    Version remoteVersion = {major, minor, patch};

    Version localVersion;
    sscanf(currentVersion, "%d.%d.%d", &localVersion.major, &localVersion.minor, &localVersion.patch);
    
    // Check if an update is available
    if (remoteVersion < localVersion) {
        ESP_LOGI(TAG, "No updates available");
        return false;
    }

    ESP_LOGD(TAG, "Manifest content: %s", payload.c_str());

    return true;
}

bool AppUpdater::updateFile(const char* remotePath, const char* localPath, const char* expectedMd5) {
    updateProgress(UpdateStatus::DOWNLOADING, 0, localPath);
    
    // Construct full URL
    String url = String(bucketUrl) + remotePath;
    ESP_LOGD(TAG, "Downloading: %s -> %s", url.c_str(), localPath);

    String localMd5 = getLocalMD5(localPath);

    if (localMd5.equals(expectedMd5)) {
        ESP_LOGI(TAG, "File already up to date: %s", localPath);
        updateProgress(UpdateStatus::SKIPPING, 100, localPath);
        return true;
    }

    // Start the download
    HTTPClient http;
    http.begin(url);
    int httpCode = http.GET();
    if (httpCode != HTTP_CODE_OK) {
        ESP_LOGE(TAG, "Download failed: %d", httpCode);
        updateProgress(UpdateStatus::ERROR, 0);
        http.end();
        return false;
    }

    // Get the stream and content length
    WiFiClient* stream = http.getStreamPtr();
    size_t contentLength = http.getSize();
    
    // Verify and save the file
    bool success = verifyAndSaveFile(stream, contentLength, localPath, expectedMd5);
    http.end();

    updateProgress(success ? UpdateStatus::COMPLETE : UpdateStatus::ERROR, 100, localPath);
    return success;
}

bool AppUpdater::verifyAndSaveFile(WiFiClient* stream, size_t contentLength, const char* localPath, const char* expectedMd5) 
{
    MD5Builder md5;
    md5.begin();
    size_t totalRead = 0;
    
    // Create temporary filename
    String tempPath = String(localPath) + ".tmp";
    
    // Open temporary file for writing
    File file = fileSystem.open(tempPath.c_str(), FILE_WRITE);
    if (!file) {
        ESP_LOGE(TAG, "Failed to open temporary file for writing");
        return false;
    }

    updateProgress(UpdateStatus::DOWNLOADING, 0, localPath);

    // Single pass: Save file and calculate MD5
    while (totalRead < contentLength) {
        size_t available = stream->available();
        if (available) {
            size_t readLen = stream->readBytes(downloadBuffer.get(), std::min(available, size_t(BUFFER_SIZE)));
            
            // Write to temp file and update MD5
            if (file.write(downloadBuffer.get(), readLen) != readLen) {
                ESP_LOGE(TAG, "Failed to write to temporary file");

                file.close();
                fileSystem.remove(tempPath.c_str());
                return false;
            }
            
            md5.add(downloadBuffer.get(), readLen);
            totalRead += readLen;
            updateProgress(UpdateStatus::DOWNLOADING, (totalRead * 90) / contentLength , localPath);
        }
        yield();
    }

    file.close();
    md5.calculate();
    String calculatedMd5 = md5.toString();

    // Verify MD5 hash
    if (!calculatedMd5.equals(expectedMd5)) {
        ESP_LOGE(TAG, "MD5 mismatch for %s", localPath);
        fileSystem.remove(tempPath.c_str());
        return false;
    }

    // Replace original file with verified temp file
    if (fileSystem.exists(localPath)) {
        fileSystem.remove(localPath);
    }
    if (!fileSystem.rename(tempPath.c_str(), localPath)) {
        ESP_LOGE(TAG, "Failed to rename temporary file");
        fileSystem.remove(tempPath.c_str());
        return false;
    }

    updateProgress(UpdateStatus::COMPLETE, 100, localPath);
    return true;
}

String AppUpdater::getLocalMD5(const char* filePath){
    File file = fileSystem.open(filePath, "r");
    if(!file){
        ESP_LOGE(TAG, "Error opening %s...", filePath);
        return String();
    }

    MD5Builder md5Builder;
    md5Builder.begin();
    size_t fileSize = file.size();
    size_t totalRead = 0;
    size_t readLen = 0;
    while (totalRead < fileSize) {
        readLen = file.readBytes(reinterpret_cast<char*>(downloadBuffer.get()), std::min(fileSize - totalRead, size_t(BUFFER_SIZE)));
        md5Builder.add(downloadBuffer.get(), readLen);
        totalRead += readLen;
    }

    md5Builder.calculate();
    file.close();
    return md5Builder.toString();
}

bool AppUpdater::updateFilesArray() {
    int successCount = 0;
    int totalFiles = jsonFilesArray.size();
    ESP_LOGI(TAG, "Found %d files in manifest", totalFiles);

    // Iterate over each file entry in the manifest
    for (JsonObject file : jsonFilesArray) {
        const char* remotePath = file["remote"];
        const char* localPath = file["local"];
        const char* expectedMd5 = file["md5"];

        // Skip invalid entries
        if (!remotePath || !localPath || !expectedMd5) {
            ESP_LOGE(TAG, "Invalid file entry in manifest");
            continue;
        }

        // Attempt to update the file
        if (updateFile(remotePath, localPath, expectedMd5)) {
            successCount++;
        }
    }

    ESP_LOGI(TAG, "Manifest update complete: %d/%d files updated", successCount, totalFiles);
    return successCount == totalFiles;
}

bool AppUpdater::updateApp() {
    updateProgress(UpdateStatus::MESSAGE, 0, "Starting firmware update");

    // Check for firmware section in manifest
    if (!jsonManifest["firmware"].is<JsonObject>() || !jsonManifest["firmware"]["md5"].is<const char*>()) {
        ESP_LOGE(TAG, "Invalid firmware section in manifest");
        updateProgress(UpdateStatus::MESSAGE, 0, "Firmware: Invalid firmware section in manifest");
        return false;
    }

    // Get the firmware MD5 hash and URL
    const char* expectedMd5 = jsonManifest["firmware"]["md5"];
    String firmwareUrl = String(bucketUrl) + appName;
    
    // Download the firmware
    HTTPClient http;
    http.begin(firmwareUrl);
    int httpCode = http.GET();
    if (httpCode != HTTP_CODE_OK) {
        ESP_LOGE(TAG, "Firmware download failed: %d", httpCode);
        updateProgress(UpdateStatus::MESSAGE, 0, "Firmware: Firmware download failed");
        http.end();
        return false;
    }

    // Check available space
    size_t firmwareSize = http.getSize();
    if (!Update.begin(firmwareSize)) {
        ESP_LOGE(TAG, "Firmware: Not enough space for update");
        updateProgress(UpdateStatus::MESSAGE, 0, "Firmware: Not enough space for update");
        http.end();
        return false;
    }

    // Set up MD5 checking
    MD5Builder md5;
    md5.begin();
    
    // Download and verify firmware
    WiFiClient* stream = http.getStreamPtr();
    size_t remaining = firmwareSize;
    while (remaining > 0) {
        size_t chunk = std::min(remaining, size_t(BUFFER_SIZE));
        size_t read = stream->readBytes(downloadBuffer.get(), chunk);
        
        // Check for timeout
        if (read == 0) {
            ESP_LOGE(TAG, "Read timeout");

            Update.abort();
            http.end();
            return false;
        }
        
        // Update MD5 and write firmware
        md5.add(downloadBuffer.get(), read);
        if (Update.write(downloadBuffer.get(), read) != read) {
            ESP_LOGE(TAG, "Write failed");
            Update.abort();
            http.end();
            return false;
        }
        
        remaining -= read;
        updateProgress(UpdateStatus::DOWNLOADING, (firmwareSize - remaining) * 100 / firmwareSize, "firmware");
    }
    
    // Verify MD5
    md5.calculate();
    String calculatedMd5 = md5.toString();
    if (!calculatedMd5.equals(expectedMd5)) {
        ESP_LOGE(TAG, "MD5 mismatch. Expected: %s, Got: %s", expectedMd5, calculatedMd5.c_str());
        updateProgress(UpdateStatus::MESSAGE, 0, "Firmware: MD5 mismatch");
        Update.abort();
        http.end();
        return false;
    }

    // Finish update
    if (!Update.end()) {
        ESP_LOGE(TAG, "Update end failed");
        updateProgress(UpdateStatus::MESSAGE, 0, "Firmware: Update failed");
        http.end();
        return false;
    }

    http.end();
    return true;
}

bool AppUpdater::IsUpdateAvailable(){
    return updateAvailable;
}



//AsyncEventSource* eventSource = nullptr;
AsyncWebSocket* wsSource = nullptr;

void startFirmwareUpdateTask(AsyncWebSocket* wsSrc) {
    wsSource = wsSrc;
    if(Update_Task_Handle) {
        ESP_LOGW(TAG, "Firmware update task already running");
        return;
    }
    xTaskCreate(firmwareUpdateTask, "FirmwareUpdate", 8192, NULL, 1, &Update_Task_Handle);
}

void firmwareUpdateTask(void* parameter) {
    static const char* TAG = "UpdateTask";
    String updateJsonPath = "/system/update.json";
    AppUpdater* updater = nullptr;

    try {
        // Read and parse update.json
        File file = LittleFS.open(updateJsonPath);
        if (!file) { 
            throw std::runtime_error("Failed to open update.json"); 
        }

        JsonDocument doc;
        DeserializationError error = deserializeJson(doc, file);
        file.close();
        
        if (error) { throw std::runtime_error("Failed to parse update.json"); }

        // Get update configuration
        JsonObject jObj = doc.as<JsonObject>();
        String baseUrl = jsonConstrainString(TAG, jObj, "baseurl", "https://storage.googleapis.com/boothifier/");
        String folderName = jsonConstrainString(TAG, jObj, "folder", "latest/");
        String url = baseUrl + folderName;

        // Initialize updater
        updater = new AppUpdater(LittleFS, FIRMWARE_VERSION, url.c_str(), "update.json", "firmware.bin");
        updater->setProgressCallback(updateProgress);
        
        ESP_LOGI(TAG, "Starting update check from: %s", url.c_str());


        // Test loop
        for(int i = 0; i < 10; i++) {
            updateProgress(AppUpdater::UpdateStatus::MESSAGE, i * 10, "test");
            vTaskDelay(2000);
        }
        updateProgress(AppUpdater::UpdateStatus::MESSAGE, 100, "test complete");
        goto end;

        // Check and perform updates
        if (!updater->checkManifest()) { throw std::runtime_error("Failed to check manifest"); }

        if (updater->IsUpdateAvailable()) {
            ESP_LOGI(TAG, "Update available, updating files...");
            
            if (!updater->updateFilesArray()) { throw std::runtime_error("Failed to update files"); }

            ESP_LOGI(TAG, "Updating firmware...");
            if (!updater->updateApp()) { throw std::runtime_error("Failed to update firmware"); }

            ESP_LOGI(TAG, "Update successful, restarting...");
            delete updater;
            vTaskDelay(2000);
            ESP.restart();
        }
        
    } catch (const std::exception& e) {
        ESP_LOGE(TAG, "Update failed: %s", e.what());
    }
    end:
    delete updater;
    Update_Task_Handle = NULL;
    vTaskDelete(NULL);
}

void updateProgress(AppUpdater::UpdateStatus newStatus, int percentage, const char* message = nullptr) {

    char buffer[128];
    const char* msg;
    bool isComplete = false;

    switch (newStatus) {
        case AppUpdater::UpdateStatus::IDLE:
            snprintf(buffer, sizeof(buffer), "Update idle");
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::MESSAGE:
            msg = message ? message : "";
            break;
        case AppUpdater::UpdateStatus::DOWNLOADING:
            snprintf(buffer, sizeof(buffer), "%s: Download progress: %d%%", message, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::VERIFYING:
            snprintf(buffer, sizeof(buffer), "%s: Verifying update: %d%%", message, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::SKIPPING:
            snprintf(buffer, sizeof(buffer), "%s: Skipping file update: already up to date", message);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::SAVING:
            snprintf(buffer, sizeof(buffer), "%s: Saving update: %d%%", message, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::COMPLETE:
            snprintf(buffer, sizeof(buffer), "%s: Update complete", message);
            msg = buffer;
            isComplete = true;
            break;
        case AppUpdater::UpdateStatus::ERROR:
            snprintf(buffer, sizeof(buffer), "%s: Update error occurred", message);
            msg = buffer;
            break;
        default:
            snprintf(buffer, sizeof(buffer), "Unknown update status: %d", (int)newStatus);
            msg = buffer;
            break;
    }

    ESP_LOGI(TAG, "%s", msg);
    sendUpdateMessage(msg, isComplete, percentage);
}

void sendUpdateMessage(const char* message, bool complete, int progress = -1) {
    if(wsSource && wsSource->count() > 0) {
        JsonDocument jsonDoc;
        jsonDoc["message"] = message;
        jsonDoc["complete"] = complete;
        jsonDoc["progress"] = progress;
        String strMessage;
        serializeJson(jsonDoc, strMessage);
        wsSource->textAll(strMessage);
    }
}






/*
void setup() {
    Serial.begin(115200);
    
    // Initialize WiFi connection first
    // ... WiFi connection code ...

    // Initialize filesystem
    if(!LittleFS.begin()) {
        Serial.println("LittleFS Mount Failed");
        return;
    }

    // Create updater instance with:
    // - Current version: "1.0.0"
    // - Update server URL: "https://my-update-server.com/"
    // - Filesystem: LittleFS
    AppUpdater updater("1.0.0", "https://storage.googleapis.com/boothifier/latest/", LittleFS);

    // Set progress callback
    updater.setProgressCallback([](int progress) {
        Serial.printf("Update progress: %d%%\n", progress);
    });

    // Check and update firmware
    if (updater.checkAndUpdate()) {
        Serial.println("Update successful! Rebooting...");
        ESP.restart();
    }

    // Update specific files from manifest
    int updatedFiles = updater.updateFilesFromManifest("test_update.json");
    Serial.printf("Updated %d files\n", updatedFiles);
}

*/