esp-recuperator/main/console/commands/cmd_wifi.c

//
// Created by MightyPork on 2018/12/08.
//

#include <freertos/FreeRTOS.h>
#include <freertos/event_groups.h>
#include <esp_wifi.h>
#include <settings.h>

#include "console/cmd_common.h"
#include <console/cmddef.h>
#include <application.h>
#include <console/prefix_match.h>


static int cmd_disable(console_ctx_t *ctx, cmd_signature_t *reg)
{
    EMPTY_CMD_SETUP("Disable WiFi");
    console_printf("WiFi "MSG_DISABLED"\nRestart to apply.\n");
    gSettings.wifi_enabled = false;
    settings_persist(SETTINGS_wifi_enabled);
    return 0;
}

static int cmd_enable(console_ctx_t *ctx, cmd_signature_t *reg)
{
    EMPTY_CMD_SETUP("Enable WiFi");
    console_printf("WiFi "MSG_ENABLED"\nRestart to apply.\n");
    gSettings.wifi_enabled = true;
    settings_persist(SETTINGS_wifi_enabled);
    return 0;
}

static const char *en_dis_cmds[] = {
        [0] = "disable",
        [1] = "enable",
        NULL
};

static int cmd_ap_conf(console_ctx_t *ctx, cmd_signature_t *reg)
{
    static struct {
        struct arg_str *cmd;
        struct arg_str *ssid;
        struct arg_str *pw;
        struct arg_str *ip;
        struct arg_end *end;
    } args;

    if (reg) {
        args.cmd = arg_str0(NULL, NULL, "{enable|disable}", EXPENDABLE_STRING("Command"));
        args.ssid = arg_str0("s", NULL, "<SSID>", EXPENDABLE_STRING("Set AP SSID"));
        args.pw = arg_str0("p", NULL, "<PWD>", EXPENDABLE_STRING("Set AP WPA2 password. Empty for open."));
        args.ip = arg_str0("a", NULL, "<IP>", "Set IP address (server + gateway). Always /24");
        args.end = arg_end(4);

        reg->argtable = &args;
        reg->help = EXPENDABLE_STRING("Configure WiFi AP mode");
        return CONSOLE_OK;
    }

    if (!g_State.wifi_inited) {
        console_printf("\x1b[31;1mWiFi is disabled!\x1b[22m\nEnable with `wifi enable`, restart to apply.\x1b[m\n");
        return 0;
    }

    if (args.cmd->count) {
        int match = prefix_match(args.cmd->sval[0], en_dis_cmds, 0);

        switch (match) {
            case 0:
                console_printf("AP mode "MSG_DISABLED"\nRestart to apply.\n");
                gSettings.ap_enabled = false;
                settings_persist(SETTINGS_ap_enabled);
                break;

            case 1:
                console_printf("AP mode "MSG_ENABLED"\nRestart to apply.\n");
                gSettings.ap_enabled = true;
                settings_persist(SETTINGS_ap_enabled);
                break;

            default:
                return CONSOLE_ERR_INVALID_ARG;
        }
    } else {
        // No cmd
        console_printf("AP mode: %s\n", gSettings.ap_enabled ? MSG_ENABLED : MSG_DISABLED);
    }

    if (args.ip->count) {
        uint32_t a = 0;
        if (!inet_aton(args.ip->sval[0], &a)) {
            console_println("Invalid IP");
            return CONSOLE_ERR_INVALID_ARG;
        }
        gSettings.ap_ip = a; // aton output is already in network byte order
        settings_persist(SETTINGS_ap_ip);

        console_println("AP IP changed, restart to apply.\n");
    }

    bool changed = false;
    wifi_config_t apconf = {};
    ESP_ERROR_CHECK(esp_wifi_get_config(ESP_IF_WIFI_AP, &apconf));

    if (args.ssid->count) {
        //apconf.ap.authmode = WIFI_AUTH_OPEN;
        strcpy((char*)apconf.ap.ssid, args.ssid->sval[0]);
        apconf.ap.ssid_len = strlen(args.ssid->sval[0]);
        changed = true;
    }

    if (args.pw->count) {
        size_t len = strlen(args.pw->sval[0]);
        if (len < 8 && len != 0) {
            console_println("AP pw must be 8 chars or more!");
            return CONSOLE_ERR_INVALID_ARG;
        }

        strcpy((char*)apconf.ap.password, args.pw->sval[0]);
        if (len == 0) {
            // if no pw is set, the AP will be open
            apconf.ap.authmode = WIFI_AUTH_OPEN;
        } else {
            apconf.ap.authmode = WIFI_AUTH_WPA2_PSK;
        }
        changed = true;
    }

    if (changed) {
        esp_err_t rv = esp_wifi_set_config(ESP_IF_WIFI_AP, &apconf);
        if (rv != ESP_OK) {
            console_printf("Error set config: %s\n", esp_err_to_name(rv));
            return -1;
        }
    }

    console_printf("AP SSID: \"%s\"\n", apconf.ap.ssid);
    console_printf("AP PW: \"%s\"\n", apconf.ap.password);
    console_printf("AP own IP: %s/24\n", inet_ntoa(gSettings.ap_ip));
    return 0;
}

static int cmd_sta_conf(console_ctx_t *ctx, cmd_signature_t *reg)
{
    static struct {
        struct arg_str *cmd;
        struct arg_end *end;
    } args;

    if (reg) {
        args.cmd = arg_str0(NULL, NULL, "{enable|disable}", EXPENDABLE_STRING("Command"));
        args.end = arg_end(1);

        reg->argtable = &args;
        reg->help = EXPENDABLE_STRING("Configure WiFi STA mode");
        return CONSOLE_OK;
    }

    if (!g_State.wifi_inited) {
        console_printf("\x1b[31;1mWiFi is disabled!\x1b[22m\nEnable with `wifi enable`, restart to apply.\x1b[m\n");
        return 0;
    }

    if (args.cmd->count) {
        int match = prefix_match(args.cmd->sval[0], en_dis_cmds, 0);

        switch (match) {
            case 0:
                console_printf("STA mode "MSG_DISABLED"\nRestart to apply.\n");
                gSettings.sta_enabled = false;
                settings_persist(SETTINGS_sta_enabled);
                break;

            case 1:
                console_printf("STA mode "MSG_ENABLED"\nRestart to apply.\n");
                gSettings.sta_enabled = true;
                settings_persist(SETTINGS_sta_enabled);
                break;

            default:
                return CONSOLE_ERR_INVALID_ARG;
        }
    } else {
        // No cmd
        console_printf("STA mode: %s\n", gSettings.sta_enabled ? MSG_ENABLED : MSG_DISABLED);
    }

    return 0;
}

/** Disconnect from WiFi and forget creds */
static int cmd_sta_forget(console_ctx_t *ctx, cmd_signature_t *reg)
{
    static struct {
        struct arg_end *end;
    } args;

    if (reg) {
        args.end = arg_end(1);

        reg->argtable = &args;
        reg->command = "wifi forget";
        reg->help = "Disconnect from WiFi AP and erase stored credentials";
        return 0;
    }

    console_printf("Removing saved WiFi credentials and disconnecting.\n");

    if (!g_State.wifi_inited) {
        console_printf("\x1b[31;1mWiFi is disabled!\x1b[22m\nEnable with `wifi enable`, restart to apply.\x1b[m\n");
        return 0;
    }

    wifi_config_t wificonf;
    esp_wifi_get_config(WIFI_IF_STA, &wificonf);
    wificonf.sta.ssid[0] = 0;
    wificonf.sta.password[0] = 0;
    esp_wifi_set_config(WIFI_IF_STA, &wificonf);
    esp_wifi_disconnect();

    return 0;
}

#define DEF_WIFI_TIMEOUT 10000

static bool wifi_join(const char* ssid, const char* pass, int timeout_ms)
{
    wifi_config_t wifi_config = {};
    strncpy((char*) wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid)-1);
    if (pass) {
        strncpy((char*) wifi_config.sta.password, pass, sizeof(wifi_config.sta.password) - 1);
    }

    ESP_ERROR_CHECK( esp_wifi_set_mode(WIFI_MODE_STA) );
    ESP_ERROR_CHECK( esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config) );

    xEventGroupClearBits(g_wifi_event_group, WIFI_CONNECTED_BIT|WIFI_FAIL_BIT);

    ESP_ERROR_CHECK( esp_wifi_disconnect() );
    ESP_ERROR_CHECK( esp_wifi_connect() );

    int bits = xEventGroupWaitBits(g_wifi_event_group, WIFI_CONNECTED_BIT|WIFI_FAIL_BIT,
            /* clear */ 0, /* wait for all */0, pdMS_TO_TICKS(timeout_ms));

    return (bits & EG_WIFI_CONNECTED_BIT) != 0;
}

static int cmd_join(console_ctx_t *ctx, cmd_signature_t *reg)
{
    static struct {
        struct arg_int *timeout;
        struct arg_str *ssid;
        struct arg_str *password;
        struct arg_end *end;
    } cmd_args;

    if (reg) {
        cmd_args.timeout = arg_int0("t", "timeout", "<t>", "Connection timeout, ms");
        cmd_args.ssid = arg_str1(NULL, NULL, "<ssid>", "SSID of AP");
        cmd_args.password = arg_str0(NULL, NULL, "<pass>", "PSK of AP");
        cmd_args.end = arg_end(2);

        reg->argtable = &cmd_args;
        reg->command = "wifi join";
        reg->help = "Join WiFi AP as a station";
        return 0;
    }

    if (!g_State.wifi_inited) {
        console_printf("\x1b[31;1mWiFi is disabled!\x1b[22m\n"
                       "Enable with `wifi enable`, restart to apply.\x1b[m\n");
        return 0;
    }

    console_printf("Connecting to '%s'\n", cmd_args.ssid->sval[0]);

    int tmeo = cmd_args.timeout->count ? cmd_args.timeout->ival[0] : DEF_WIFI_TIMEOUT;

    bool connected = wifi_join(cmd_args.ssid->sval[0],
                               cmd_args.password->sval[0],
                               tmeo);
    if (!connected) {
        console_printf("Connection timed out\n");

        // erase config
        wifi_config_t wifi_config = {};
        ESP_ERROR_CHECK( esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config) );

        return 1;
    }
    console_printf("Connected\n");
    return 0;
}

static int cmd_wifi_status(console_ctx_t *ctx, cmd_signature_t *reg)
{
    static struct {
        struct arg_end *end;
    } cmd_args;

    if (reg) {
        cmd_args.end = arg_end(1);

        reg->argtable = &cmd_args;
        reg->command = "wifi status";
        reg->help = "Check WiFi / IP status";
        return 0;
    }

    console_printf("WiFi support: %s\n", gSettings.wifi_enabled ? MSG_ENABLED : MSG_DISABLED);
    console_printf("STA mode: %s\n", gSettings.sta_enabled ? MSG_ENABLED : MSG_DISABLED);
    console_printf("AP mode: %s\n", gSettings.ap_enabled ? MSG_ENABLED : MSG_DISABLED);

    console_printf("\n");

    if (gSettings.wifi_enabled) {
        wifi_config_t config;
        if (ESP_OK == esp_wifi_get_config(ESP_IF_WIFI_STA, &config)) {
            if (config.sta.ssid[0]) {
                console_printf("Configured to connect to SSID \"%s\".\n"
                               "Use `wifi forget` to drop saved credentials.\n\n",
                               config.sta.ssid);

                tcpip_adapter_ip_info_t ipinfo;
                if (ESP_OK == tcpip_adapter_get_ip_info(TCPIP_ADAPTER_IF_STA, &ipinfo)) {
                    // ! inet_ntoa uses a global static buffer, cant use multiple in one printf call
                    console_printf("IP: %s, ", inet_ntoa(ipinfo.ip.addr));
                    console_printf("Mask: %s, ", inet_ntoa(ipinfo.netmask.addr));
                    console_printf("Gateway: %s\n", inet_ntoa(ipinfo.gw.addr));
                } else {
                    console_printf("No IP!\n");
                }
            }
            else {
                console_printf("No network SSID configured.\n"
                               "Use `wifi join` to connect to a WiFi network.\n");
            }
        }
    }

    return 0;
}

void register_cmd_wifi(void)
{
    console_group_add("wifi", "WiFi configuration");

    console_cmd_register(cmd_enable, "wifi enable");
    console_cmd_register(cmd_disable, "wifi disable");

    console_cmd_register(cmd_sta_conf, "wifi sta");
    console_cmd_register(cmd_sta_forget, "wifi forget");
    console_cmd_register(cmd_join, "wifi join");

    console_cmd_register(cmd_ap_conf, "wifi ap");

    console_cmd_register(cmd_wifi_status, "wifi status");
}