boothifier/temporary/AppUpgrade orig.cpp

#include "AppUpgrade.h"
#include "esp_log.h"
#include <MD5Builder.h>
#include <LittleFS.h>
#include <memory>
#include "global.h"
#include "JsonConstrain.h"
#include "BLE_UpdateService.h"
#include <HTTPClient.h>
#include <Update.h>

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

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

String updateUrl = "";

Version otaVersion;


AppUpdater::AppUpdater(fs::FS& fs, Version localVersion, const char* bucket, const char* manifestName, const char* appBin)
    : localVersion(localVersion), manifestName(manifestName), appName(appBin), fileSystem(fs), downloadBuffer(new uint8_t[BUFFER_SIZE])
{
    baseUrl = bucket ? String(bucket) : String(DEFAULT_MANIFEST_URL);
    // Ensure baseUrl ends with a single '/'
    if(!baseUrl.endsWith("/")) baseUrl += "/";
    ESP_LOGI(TAG, "AppUpdater initialized (local v%s) baseUrl=%s", localVersion.toString().c_str(), baseUrl.c_str());
}

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 = buildUrl(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);
    ESP_LOGD(TAG, "Manifest deserialized");
    if (error) {
        ESP_LOGE(TAG, "Failed to parse manifest: %s", error.c_str());
        return false;
    }

    // Check for files section
    jsonFilesArray = jsonManifest["files"];
    if (jsonFilesArray.isNull()) {
        ESP_LOGE(TAG, "No files section in manifest");
        return false;
    }else{
        ESP_LOGD(TAG, "%d Files found", jsonFilesArray.size());
    }

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

    // Get the remote version
    byte major = jsonVersion["major"] | 0;
    byte minor = jsonVersion["minor"] | 0;
    byte patch = jsonVersion["patch"] | 0;
    otaVersion = {major, minor, patch};

    //Version localVersion;
    //::sscanf(localVersion, "%d.%d.%d", &localVersion.major, &localVersion.minor, &localVersion.patch);
    
    // Check if an update is available
    updateAvailable = false;
    // Only mark update available if remote is strictly newer than local
    if (otaVersion <= localVersion) {
        ESP_LOGI(TAG, "No updates available");
        return false;
    }else{
        updateAvailable = true;
        ESP_LOGD(TAG, "Update available");
    }

    //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 = buildUrl(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::FILE_SKIPPED, 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, "Download failed");
        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();
    if(!success){
        updateProgress( UpdateStatus::ERROR, 0, "MD5 verification failed");
    }else{
        updateProgress( UpdateStatus::FILE_SAVED, 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);

    if (contentLength > 0) {
        // Single pass with known content length
        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();
        }
    } else {
        // Unknown content length: read until stream ends
        for (;;) {
            size_t readLen = stream->readBytes(downloadBuffer.get(), BUFFER_SIZE);
            if (readLen == 0) {
                break;
            }
            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;
            // Progress unknown; emit periodic heartbeats at 0%
            updateProgress(UpdateStatus::DOWNLOADING, 0, 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;
    }

    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::ERROR, 0, "Firmware: Invalid firmware section in manifest");
        return false;
    }

    // Get the firmware MD5 hash and URL
    const char* expectedMd5 = jsonManifest["firmware"]["md5"];
    String firmwareUrl = buildUrl(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::ERROR, 0, "Firmware: Firmware download failed");
        http.end();
        return false;
    }

    // Check available space
    size_t firmwareSize = http.getSize();
    if (!Update.begin(firmwareSize > 0 ? firmwareSize : UPDATE_SIZE_UNKNOWN)) {
        ESP_LOGE(TAG, "Firmware: Not enough space for update");
        updateProgress(UpdateStatus::ERROR, 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();
    if (firmwareSize > 0) {
        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");
        }
    } else {
        // Unknown size: stream until end
        for (;;) {
            size_t read = stream->readBytes(downloadBuffer.get(), BUFFER_SIZE);
            if (read == 0) break;
            md5.add(downloadBuffer.get(), read);
            if (Update.write(downloadBuffer.get(), read) != read) {
                ESP_LOGE(TAG, "Write failed");
                Update.abort();
                http.end();
                return false;
            }
            updateProgress(UpdateStatus::DOWNLOADING, 0, "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::MD5_FAILED, 0, "Firmware: MD5 mismatch");
        Update.abort();
        http.end();
        return false;
    }

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

    http.end();
    updateProgress(UpdateStatus::COMPLETE, 0, "Firmware: Complete");
    return true;
}

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

String AppUpdater::buildUrl(const char* path) const {
    if(!path || !*path) return baseUrl; // just base
    String p(path);
    // If already absolute URL, pass through
    if(p.startsWith("http://") || p.startsWith("https://")) return p;
    // Strip leading slashes to avoid double
    while(p.startsWith("/")) p.remove(0,1);
    // Ensure baseUrl has single trailing slash
    String b = baseUrl;
    if(!b.endsWith("/")) b += "/";
    return b + p;
}


AsyncEventSource* eventProgress = nullptr;
void startFirmwareUpdateTask(AsyncEventSource* evProg) {
    eventProgress = evProg;
    if(Update_Task_Handle) {
        ESP_LOGW(TAG, "Firmware update task already running");
        return;
    }
    xTaskCreate(firmwareUpdateTask, "FirmwareUpdate", 1024*8, NULL, 1, &Update_Task_Handle);
}

void firmwareUpdateTask(void* parameter) {
    static const char* TAG = "UpdateTask";
    AppUpdater* updater = nullptr;

    try {
        loadUpdateJson();

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

        // 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...");
            
            sendUpdateMessage("Restarting ", true, 100);
            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 startVersionCheckTask() {
    if(versionCheckTask_Handle != NULL) {
        ESP_LOGW(TAG, "Version Check Tak already running");
        return;
    }
    xTaskCreate(versionCheckTask, "VersionCheckTask", 1024*8, NULL, 1, &versionCheckTask_Handle);
}

void versionCheckTask(void* parameter){

    if(updateUrl == ""){
        loadUpdateJson();
    }

    if(checkManifest(otaVersion) == false){
        ESP_LOGE(TAG, "Error checking manifest");
    }

    versionCheckTask_Handle = NULL;
    vTaskDelete(NULL);
}

void loadUpdateJson(void) {
    try {
        ESP_LOGD(TAG, "loadUpdateJaon function...");
        if(updateUrl == "") {
            String updateJsonPath = "/system/update.json";

            // 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 folderName = jsonConstrainString(TAG, jObj, "folder", "latest/");
            String baseUrl = jsonConstrainString(TAG, jObj, "baseurl", "https://s3-minio.boothwizard.com/boothifier/");
            updateUrl = baseUrl + folderName;

            ESP_LOGD(TAG, "updateUrl: %s", updateUrl.c_str());
        }
    } catch (const std::exception& e) {
        ESP_LOGE(TAG, "Update failed: %s", e.what());
    }
}

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

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

    const char* safeMsg = message ? message : "";
    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%%", safeMsg, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::VERIFYING:
            snprintf(buffer, sizeof(buffer), "%s: Verifying update: %d%%", safeMsg, percentage);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::FILE_SKIPPED:
            snprintf(buffer, sizeof(buffer), "%s: Skipping file update, already up to date", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::FILE_SAVED:
            snprintf(buffer, sizeof(buffer), "%s: File Saved", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::MD5_FAILED:
            snprintf(buffer, sizeof(buffer), "%s: MD5 Verification Failed", safeMsg);
            msg = buffer;
            break;
        case AppUpdater::UpdateStatus::COMPLETE:
            snprintf(buffer, sizeof(buffer), "Firmware Update Complete!!!");
            msg = buffer;
            isComplete = true;
            break;
        case AppUpdater::UpdateStatus::ERROR:
            snprintf(buffer, sizeof(buffer), "Error!: %s", safeMsg);
            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(eventProgress && eventProgress->count() > 0) {
        JsonDocument jsonDoc;
        jsonDoc["message"] = message;
        jsonDoc["complete"] = complete;
        jsonDoc["progress"] = progress;
        String strMessage;
        serializeJson(jsonDoc, strMessage);
        eventProgress->send(strMessage.c_str(), "update", millis());
    }
    else{
        ESP_LOGW(TAG, "No clients connected to event source");
    }

    bleUpgrade_send_message(message);
}

bool checkManifest(Version& remoteVersion) {
    const char* TAG = "manifestCheck";
    String url = updateUrl + "update.json";
    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();

    //ESP_LOGD(TAG, "%s", payload.c_str());

    // Parse JSON
    JsonDocument jsonManifest;
    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
    byte major = jsonVersion["major"] | 0;
    byte minor = jsonVersion["minor"] | 0;
    byte patch = jsonVersion["patch"] | 0;
    remoteVersion = {major, minor, patch};
    ESP_LOGI(TAG, "Remote version: %s", remoteVersion.toString().c_str());
    return true;
}

/*
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);
}

*/