knxdisplay/main/Wifi.cpp

#include "Wifi.hpp"
#include "sdkconfig.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "nvs_flash.h"
#include "nvs.h"
#include <cstring>
#include <algorithm>

static const char* TAG = "Wifi";

static const char* disconnectReasonToString(uint16_t reason) {
    switch (reason) {
        case WIFI_REASON_UNSPECIFIED:
            return "UNSPECIFIED";
        case WIFI_REASON_AUTH_EXPIRE:
            return "AUTH_EXPIRE";
        case WIFI_REASON_AUTH_LEAVE:
            return "AUTH_LEAVE";
        case WIFI_REASON_ASSOC_EXPIRE:
            return "ASSOC_EXPIRE";
        case WIFI_REASON_NOT_AUTHED:
            return "NOT_AUTHED";
        case WIFI_REASON_NOT_ASSOCED:
            return "NOT_ASSOCED";
        case WIFI_REASON_ASSOC_LEAVE:
            return "ASSOC_LEAVE";
        case WIFI_REASON_4WAY_HANDSHAKE_TIMEOUT:
            return "4WAY_HANDSHAKE_TIMEOUT";
        case WIFI_REASON_GROUP_KEY_UPDATE_TIMEOUT:
            return "GROUP_KEY_UPDATE_TIMEOUT";
        case WIFI_REASON_BEACON_TIMEOUT:
            return "BEACON_TIMEOUT";
        case WIFI_REASON_NO_AP_FOUND:
            return "NO_AP_FOUND";
        case WIFI_REASON_AUTH_FAIL:
            return "AUTH_FAIL";
        case WIFI_REASON_ASSOC_FAIL:
            return "ASSOC_FAIL";
        case WIFI_REASON_HANDSHAKE_TIMEOUT:
            return "HANDSHAKE_TIMEOUT";
        case WIFI_REASON_CONNECTION_FAIL:
            return "CONNECTION_FAIL";
        default:
            return "UNKNOWN";
    }
}

static bool isWpa3AuthMode(wifi_auth_mode_t authmode) {
    switch (authmode) {
        case WIFI_AUTH_WPA3_PSK:
        case WIFI_AUTH_WPA3_EXT_PSK:
        case WIFI_AUTH_WPA3_EXT_PSK_MIXED_MODE:
            return true;
        default:
            return false;
    }
}

Wifi& Wifi::instance() {
    static Wifi instance;
    return instance;
}

void Wifi::eventHandler(void* arg, esp_event_base_t event_base,
                        int32_t event_id, void* event_data) {
    Wifi& wifi = Wifi::instance();

    if (event_base == WIFI_EVENT) {
        switch (event_id) {
            case WIFI_EVENT_STA_START:
                ESP_LOGI(TAG, "WiFi station started");
                break;
            case WIFI_EVENT_STA_DISCONNECTED:
                if (event_data) {
                    wifi_event_sta_disconnected_t* event =
                        (wifi_event_sta_disconnected_t*)event_data;
                    ESP_LOGI(TAG, "WiFi disconnected (reason=%u/%s/%s)",
                             event->reason, disconnectReasonToString(event->reason), event->ssid);
                } else {
                    ESP_LOGI(TAG, "WiFi disconnected (reason=unknown)");
                }
                wifi.connected_ = false;
                wifi.currentSSID_ = "";
                if (wifi.statusCallback_) {
                    wifi.statusCallback_(false);
                }

                xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_FAIL_BIT);
                
                break;
            case WIFI_EVENT_SCAN_DONE:
                ESP_LOGI(TAG, "WiFi scan done");
                xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_SCAN_DONE_BIT);
                break;
            default:
                break;
        }
    } else if (event_base == IP_EVENT) {
        if (event_id == IP_EVENT_STA_GOT_IP) {
            ip_event_got_ip_t* event = (ip_event_got_ip_t*)event_data;
            ESP_LOGI(TAG, "Got IP: " IPSTR, IP2STR(&event->ip_info.ip));
            wifi.connected_ = true;
            if (wifi.statusCallback_) {
                wifi.statusCallback_(true);
            }
            xEventGroupSetBits(wifi.wifiEventGroup_, WIFI_CONNECTED_BIT);
        }
    }
}

void Wifi::init(bool autoConnect) {
    if (initialized_) {
        return;
    }

    ESP_LOGI(TAG, "Initializing WiFi");

    wifiEventGroup_ = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    if (!autoConnect) {
        ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
        ESP_LOGI(TAG, "WiFi storage set to RAM (auto-connect disabled)");
    }

    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &eventHandler,
                                                        nullptr,
                                                        nullptr));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &eventHandler,
                                                        nullptr,
                                                        nullptr));

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    int8_t tx_power = 40; // 0.25 dBm units -> 5 dBm.
    esp_err_t tx_err = esp_wifi_set_max_tx_power(tx_power);
    if (tx_err != ESP_OK) {
        ESP_LOGW(TAG, "Failed to set max TX power to %d: %s", tx_power, esp_err_to_name(tx_err));
    } else {
        ESP_LOGI(TAG, "Set max TX power to %d (0.25 dBm units)", tx_power);
    }

    initialized_ = true;

    if (autoConnect) {
        auto savedNetworks = getSavedNetworks();
        if (!savedNetworks.empty()) {
            std::string password;
            if (getSavedPassword(savedNetworks[0].c_str(), password)) {
                ESP_LOGI(TAG, "Auto-connecting to saved network: %s", savedNetworks[0].c_str());
                connect(savedNetworks[0].c_str(), password.c_str());
            }
        }
    }

    ESP_LOGI(TAG, "WiFi initialized");
}

void Wifi::scan(std::function<void(std::vector<wifi_ap_record_t>&)> callback) {
    scanCallback_ = callback;

    xEventGroupClearBits(wifiEventGroup_, WIFI_SCAN_DONE_BIT);

    wifi_scan_config_t scan_config = {};
    scan_config.ssid = nullptr;
    scan_config.bssid = nullptr;
    scan_config.channel = 0;
    scan_config.show_hidden = false;
    scan_config.scan_type = WIFI_SCAN_TYPE_ACTIVE;
    scan_config.scan_time.active.min = 100;
    scan_config.scan_time.active.max = 300;

    ESP_ERROR_CHECK(esp_wifi_scan_start(&scan_config, false));

    xTaskCreate([](void* arg) {
        Wifi& wifi = Wifi::instance();

        xEventGroupWaitBits(wifi.wifiEventGroup_, WIFI_SCAN_DONE_BIT,
                           pdTRUE, pdFALSE, pdMS_TO_TICKS(10000));

        uint16_t ap_count = 0;
        esp_wifi_scan_get_ap_num(&ap_count);

        std::vector<wifi_ap_record_t> ap_records(ap_count);
        if (ap_count > 0) {
            esp_wifi_scan_get_ap_records(&ap_count, ap_records.data());
        }

        ESP_LOGI(TAG, "Found %d access points", ap_count);

        if (wifi.scanCallback_) {
            wifi.scanCallback_(ap_records);
        }

        vTaskDelete(nullptr);
    }, "wifi_scan", 4096, nullptr, 5, nullptr);
}

void Wifi::connect(const char* ssid, const char* password, wifi_auth_mode_t authmode) {
    ESP_LOGI(TAG, "Connecting to %s with password: '%s' (authmode=%d)", ssid, password, authmode);

    xEventGroupClearBits(wifiEventGroup_, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT);

    wifi_config_t wifi_config = wifi_config_t();
    wifi_sta_config_t sta_config = wifi_sta_config_t();
    wifi_config.sta = sta_config;

    strcpy((char *)wifi_config.sta.ssid, (char *)ssid);         
    strcpy((char *)wifi_config.sta.password, (char *)password); 
    
    wifi_config.sta.pmf_cfg.capable = true;
    wifi_config.sta.pmf_cfg.required = false;

    

    currentSSID_ = ssid;

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));

    esp_err_t err = esp_wifi_set_config(WIFI_IF_STA, &wifi_config);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to set WiFi config: %s", esp_err_to_name(err));
        return;
    }

    err = esp_wifi_connect();
    if (err == ESP_ERR_WIFI_CONN) {
        ESP_LOGW(TAG, "WiFi connect failed: ESP_ERR_WIFI_CONN");
        return;
    }
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "WiFi connect failed: %s", esp_err_to_name(err));
        return;
    }
}

void Wifi::disconnect() {
    ESP_LOGI(TAG, "Disconnecting");
    esp_wifi_disconnect();
    connected_ = false;
    currentSSID_ = "";
}

bool Wifi::isConnected() {
    return connected_;
}

std::string Wifi::getCurrentSSID() {
    return currentSSID_;
}

std::string Wifi::getIPAddress() {
    esp_netif_t* netif = esp_netif_get_handle_from_ifkey("WIFI_STA_DEF");
    if (!netif) {
        return "";
    }

    esp_netif_ip_info_t ip_info;
    if (esp_netif_get_ip_info(netif, &ip_info) != ESP_OK) {
        return "";
    }

    char ip_str[16];
    snprintf(ip_str, sizeof(ip_str), IPSTR, IP2STR(&ip_info.ip));
    return std::string(ip_str);
}

int8_t Wifi::getRSSI() {
    wifi_ap_record_t ap_info;
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        return ap_info.rssi;
    }
    return 0;
}

void Wifi::setStatusCallback(std::function<void(bool connected)> callback) {
    statusCallback_ = callback;
}

void Wifi::saveNetwork(const char* ssid, const char* password) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error opening NVS: %s", esp_err_to_name(err));
        return;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);

    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }

    err = nvs_set_str(handle, key, password);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error saving password: %s", esp_err_to_name(err));
    }

    auto networks = getSavedNetworks();
    bool found = false;
    for (const auto& net : networks) {
        if (net == ssid) {
            found = true;
            break;
        }
    }

    if (!found) {
        networks.insert(networks.begin(), ssid);
        if (networks.size() > MAX_SAVED_NETWORKS) {
            networks.pop_back();
        }

        std::string networkList;
        for (size_t i = 0; i < networks.size(); i++) {
            if (i > 0) networkList += ";";
            networkList += networks[i];
        }
        nvs_set_str(handle, "networks", networkList.c_str());
    }

    nvs_commit(handle);
    nvs_close(handle);

    ESP_LOGI(TAG, "Saved network: %s", ssid);
}

void Wifi::removeNetwork(const char* ssid) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Error opening NVS: %s", esp_err_to_name(err));
        return;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);
    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }
    nvs_erase_key(handle, key);

    auto networks = getSavedNetworks();
    std::string ssidStr(ssid);
    networks.erase(std::remove(networks.begin(), networks.end(), ssidStr), networks.end());

    std::string networkList;
    for (size_t i = 0; i < networks.size(); i++) {
        if (i > 0) networkList += ";";
        networkList += networks[i];
    }
    nvs_set_str(handle, "networks", networkList.c_str());

    nvs_commit(handle);
    nvs_close(handle);

    ESP_LOGI(TAG, "Removed network: %s", ssid);
}

std::vector<std::string> Wifi::getSavedNetworks() {
    std::vector<std::string> networks;

    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return networks;
    }

    size_t required_size = 0;
    err = nvs_get_str(handle, "networks", nullptr, &required_size);
    if (err == ESP_OK && required_size > 1) {
        char* buf = new char[required_size];
        nvs_get_str(handle, "networks", buf, &required_size);
        std::string networkList(buf);
        delete[] buf;

        if (!networkList.empty()) {
            size_t start = 0;
            size_t end = networkList.find(';');
            while (end != std::string::npos) {
                std::string net = networkList.substr(start, end - start);
                if (!net.empty()) {
                    networks.push_back(net);
                }
                start = end + 1;
                end = networkList.find(';', start);
            }
            if (start < networkList.size()) {
                std::string net = networkList.substr(start);
                if (!net.empty()) {
                    networks.push_back(net);
                }
            }
        }
    }

    nvs_close(handle);
    return networks;
}

bool Wifi::getSavedPassword(const char* ssid, std::string& password) {
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return false;
    }

    char key[32];
    snprintf(key, sizeof(key), "pw_%s", ssid);
    size_t key_len = strlen(key);
    if (key_len > 15) {
        key[15] = '\0';
    }

    size_t required_size = 0;
    err = nvs_get_str(handle, key, nullptr, &required_size);
    if (err == ESP_OK && required_size > 0) {
        password.resize(required_size);
        nvs_get_str(handle, key, &password[0], &required_size);
        nvs_close(handle);
        return true;
    }

    nvs_close(handle);
    return false;
}