Chapter 40: WiFi WPS Implementation

// main.c
#include <stdio.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "esp_netif.h"

static const char *TAG = "WPS_EXAMPLE";

// Assume s_wifi_event_group and bits are defined elsewhere
extern EventGroupHandle_t s_wifi_event_group;
extern const int WIFI_CONNECTED_BIT;
extern const int WIFI_FAIL_BIT; // Reuse for general failure/timeout

// --- WPS Configuration ---
// Define the WPS method we want to use
#define ESP_WPS_MODE WPS_TYPE_PBC // Or WPS_TYPE_PIN

// Forward declarations
static void wifi_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data);
static void start_wps(void);

// Modified WiFi Event Handler
static void wifi_event_handler(void* arg, esp_event_base_t event_base,
                               int32_t event_id, void* event_data)
{
    if (event_base == WIFI_EVENT) {
        switch (event_id) {
            case WIFI_EVENT_STA_START:
                ESP_LOGI(TAG, "WIFI_EVENT_STA_START");
                // Optionally start WPS immediately, or trigger via button press
                // start_wps(); // Example: Trigger WPS start here
                break;
            case WIFI_EVENT_STA_CONNECTED:
                 ESP_LOGI(TAG, "WIFI_EVENT_STA_CONNECTED - Should not happen during WPS");
                 // If WPS was successful, it triggers WIFI_EVENT_STA_WPS_ER_SUCCESS first,
                 // then the stack uses the obtained credentials to connect, leading to
                 // normal WIFI_EVENT_STA_CONNECTED and IP_EVENT_STA_GOT_IP events later.
                break;
            case WIFI_EVENT_STA_DISCONNECTED:
                ESP_LOGW(TAG, "WIFI_EVENT_STA_DISCONNECTED");
                // Could happen if connection attempt after successful WPS fails
                // Or if connection drops after successful WPS connection
                is_connected = false; // Assuming global is_connected flag
                if (s_wifi_event_group != NULL) {
                    xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT);
                    xEventGroupClearBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
                }
                // Optional: Attempt reconnect or restart WPS?
                // esp_wifi_connect(); // Simple reconnect attempt
                break;

            // --- WPS Specific Events ---
            case WIFI_EVENT_STA_WPS_ER_SUCCESS:
                ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_SUCCESS: WPS Success! Credentials obtained.");
                // WPS finished successfully. The ESP-IDF WiFi stack should now
                // automatically try to connect using the obtained credentials.
                // We just need to disable the WPS mode.
                ESP_ERROR_CHECK(esp_wifi_wps_disable());
                ESP_LOGI(TAG, "WPS Disabled. Waiting for connection...");
                // Now wait for WIFI_EVENT_STA_CONNECTED and IP_EVENT_STA_GOT_IP
                break;
            case WIFI_EVENT_STA_WPS_ER_FAILED:
                ESP_LOGE(TAG, "WIFI_EVENT_STA_WPS_ER_FAILED: WPS Failed.");
                ESP_ERROR_CHECK(esp_wifi_wps_disable()); // Disable WPS on failure
                if (s_wifi_event_group != NULL) {
                    xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT); // Signal failure
                }
                // Retry WPS or enter fallback mode?
                break;
            case WIFI_EVENT_STA_WPS_ER_TIMEOUT:
                ESP_LOGW(TAG, "WIFI_EVENT_STA_WPS_ER_TIMEOUT: WPS Timeout.");
                ESP_ERROR_CHECK(esp_wifi_wps_disable()); // Disable WPS on timeout
                if (s_wifi_event_group != NULL) {
                    xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT); // Signal failure
                }
                // Retry WPS or enter fallback mode?
                break;
            case WIFI_EVENT_STA_WPS_ER_PIN:
                 ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_PIN event received");
                 // Handle PIN logic if using PIN mode (e.g., display AP PIN)
                 // For PBC mode, this event is usually ignored.
                break;
            default:
                 ESP_LOGD(TAG, "Unhandled WIFI_EVENT: %ld", event_id);
                break;
        }
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
        ESP_LOGI(TAG, "IP_EVENT_STA_GOT_IP: Got IP:" IPSTR, IP2STR(&event->ip_info.ip));
        is_connected = true; // Assuming global is_connected flag
         if (s_wifi_event_group != NULL) {
            xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
            xEventGroupClearBits(s_wifi_event_group, WIFI_FAIL_BIT);
         }
        // Connection successful (likely after WPS success)
    } else {
         ESP_LOGD(TAG, "Received unhandled event: base=%s, id=%ld", event_base, event_id);
    }
}

// main.c

// Function to configure and start WPS
static void start_wps(void) {
    ESP_LOGI(TAG, "Starting WPS (%s mode)...", (ESP_WPS_MODE == WPS_TYPE_PBC) ? "PBC" : "PIN");

    // Configure WPS
    wifi_wps_config_t config = WIFI_WPS_CONFIG_PBC_SYSTEM_INIT; // Helper macro for PBC defaults
    // If using PIN, use WIFI_WPS_CONFIG_PIN_SYSTEM_INIT and set config.pin

    // Override defaults if needed:
    config.wps_type = ESP_WPS_MODE;
    // config.factory_info.manufacturer = "ESPRESSIF"; // Optional
    // config.factory_info.model_number = "ESP32";     // Optional
    // config.factory_info.model_name = "ESPRESSIF IOT"; // Optional
    // config.factory_info.device_name = "ESP STATION"; // Optional

    // Enable WPS
    esp_err_t err = esp_wifi_wps_enable(&config);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to enable WPS: %s", esp_err_to_name(err));
        return; // Cannot start if enable failed
    }
    ESP_LOGI(TAG, "WPS Enabled.");

    // Start WPS (timeout handled by events)
    err = esp_wifi_wps_start(0); // 0 -> indefinite wait (use events for timeout)
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to start WPS: %s", esp_err_to_name(err));
        esp_wifi_wps_disable(); // Disable if start failed
        return;
    }
    ESP_LOGI(TAG, "WPS Started. Please activate WPS on your router (PBC mode).");
     if (s_wifi_event_group != NULL) {
         // Clear any previous failure bits before starting
         xEventGroupClearBits(s_wifi_event_group, WIFI_FAIL_BIT);
     }
}

// --- Example: Trigger WPS from a task or button handler ---
// This function simulates a button press triggering WPS
void trigger_wps_start(void) {
     // Ensure WiFi station is started before initiating WPS
     // (Assuming wifi_init_sta was called earlier)
     start_wps();
}

// main.c

// Event group handle defined globally or passed appropriately
EventGroupHandle_t s_wifi_event_group;
const int WIFI_CONNECTED_BIT = BIT0;
const int WIFI_FAIL_BIT = BIT1;
bool is_connected = false; // Simple global state

// Basic WiFi Init (registers event handler)
static void wifi_init(void) {
    s_wifi_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_t *sta_netif = esp_netif_create_default_wifi_sta();
    assert(sta_netif); // Ensure netif was created

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // Register event handlers for ALL WiFi and IP events
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &wifi_event_handler,
                                                        NULL,
                                                        NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &wifi_event_handler,
                                                        NULL,
                                                        NULL));

    // Set mode to Station
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    // Do NOT set specific SSID/Password here if using WPS for initial setup
    // Do NOT call esp_wifi_connect() here if triggering WPS manually

    // Start the WiFi driver tasks
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG, "WiFi Initialized. Ready for WPS.");
}


void app_main(void)
{
    // Initialize NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
      ESP_ERROR_CHECK(nvs_flash_erase());
      ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    ESP_LOGI(TAG, "ESP32 WPS Example");

    // Initialize WiFi stack and event handlers
    wifi_init();

    // --- Trigger WPS ---
    // For testing, trigger WPS after a short delay.
    // In a real application, trigger this from a button press or other user action.
    vTaskDelay(pdMS_TO_TICKS(5000)); // Wait 5 seconds
    trigger_wps_start();

    // --- Wait for Result ---
    // Example: Wait here for connection or failure using the event group
    ESP_LOGI(TAG, "Waiting for WPS/Connection result...");
    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group,
                                           WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
                                           pdFALSE, // Don't clear on exit
                                           pdFALSE, // Wait for either bit
                                           portMAX_DELAY); // Wait indefinitely

    if (bits & WIFI_CONNECTED_BIT) {
        ESP_LOGI(TAG, "Successfully connected via WPS!");
        // Proceed with application logic
    } else if (bits & WIFI_FAIL_BIT) {
        ESP_LOGE(TAG, "WPS or subsequent connection failed.");
        // Handle failure: retry, enter manual config mode, etc.
    } else {
        ESP_LOGE(TAG, "Unexpected event group result.");
    }


    // Keep main task running (or do other things)
    while(1) {
        vTaskDelay(pdMS_TO_TICKS(60000)); // Idle
    }
}