From 2e52c0ae62b53245b8f054808dd5b178837a6c99 Mon Sep 17 00:00:00 2001
From: anecdata <16617689+anecdata@users.noreply.github.com>
Date: Sat, 24 Apr 2021 16:41:41 -0500
Subject: [PATCH] set all wifi modes; add sta start & sta/ap stop

---
 ports/esp32s2/common-hal/wifi/Radio.c    | 86 +++++++++++++++++-------
 ports/esp32s2/common-hal/wifi/__init__.c |  3 +
 shared-bindings/wifi/Radio.c             | 34 +++++++++-
 shared-bindings/wifi/Radio.h             |  3 +
 4 files changed, 99 insertions(+), 27 deletions(-)

diff --git a/ports/esp32s2/common-hal/wifi/Radio.c b/ports/esp32s2/common-hal/wifi/Radio.c
index 85deda704c..76d7ce3a22 100644
--- a/ports/esp32s2/common-hal/wifi/Radio.c
+++ b/ports/esp32s2/common-hal/wifi/Radio.c
@@ -42,28 +42,44 @@
 
 #define MAC_ADDRESS_LENGTH 6
 
-static void start_station(wifi_radio_obj_t *self) {
-    if (self->sta_mode) {
-        return;
-    }
+static void set_mode_station(wifi_radio_obj_t *self, bool state) {
     wifi_mode_t next_mode;
-    if (self->ap_mode) {
-        next_mode = WIFI_MODE_APSTA;
-    } else {
-        next_mode = WIFI_MODE_STA;
-    }
-    esp_wifi_set_mode(next_mode);
+    if (state) {
+        if (self->ap_mode) {
+            next_mode = WIFI_MODE_APSTA;
+        } else {
+            next_mode = WIFI_MODE_STA;
+        }
+     } else {
+        if (self->ap_mode) {
+            next_mode = WIFI_MODE_AP;
+        } else {
+            next_mode = WIFI_MODE_NULL;
+        }
+     }
 
-    self->sta_mode = 1;
+    esp_wifi_set_mode(next_mode);
+    self->sta_mode = state;
 }
 
-static void start_ap(wifi_radio_obj_t *self) {
-    if (self->ap_mode) {
-        return;
-    }
-    esp_wifi_set_mode(WIFI_MODE_APSTA);
+static void set_mode_ap(wifi_radio_obj_t *self, bool state) {
+    wifi_mode_t next_mode;
+    if (state) {
+        if (self->sta_mode) {
+            next_mode = WIFI_MODE_APSTA;
+        } else {
+            next_mode = WIFI_MODE_AP;
+        }
+     } else {
+        if (self->sta_mode) {
+            next_mode = WIFI_MODE_STA;
+        } else {
+            next_mode = WIFI_MODE_NULL;
+        }
+     }
 
-    self->ap_mode = 1;
+    esp_wifi_set_mode(next_mode);
+    self->ap_mode = state;
 }
 
 bool common_hal_wifi_radio_get_enabled(wifi_radio_obj_t *self) {
@@ -80,8 +96,6 @@ void common_hal_wifi_radio_set_enabled(wifi_radio_obj_t *self, bool enabled) {
         return;
     }
     if (!self->started && enabled) {
-        // esp_wifi_start() would default to soft-AP, thus setting it to station
-        start_station(self);
         ESP_ERROR_CHECK(esp_wifi_start());
         self->started = true;
         return;
@@ -107,7 +121,7 @@ mp_obj_t common_hal_wifi_radio_start_scanning_networks(wifi_radio_obj_t *self) {
     if (!common_hal_wifi_radio_get_enabled(self)) {
         mp_raise_RuntimeError(translate("wifi is not enabled"));
     }
-    start_station(self);
+    set_mode_station(self, true);
 
     wifi_scannednetworks_obj_t *scan = m_new_obj(wifi_scannednetworks_obj_t);
     scan->base.type = &wifi_scannednetworks_type;
@@ -142,12 +156,28 @@ void common_hal_wifi_radio_set_hostname(wifi_radio_obj_t *self, const char *host
     esp_netif_set_hostname(self->netif, hostname);
 }
 
+void common_hal_wifi_radio_start_station(wifi_radio_obj_t *self) {
+    if (!common_hal_wifi_radio_get_enabled(self)) {
+        mp_raise_RuntimeError(translate("wifi is not enabled"));
+    }
+
+    set_mode_station(self, true);
+}
+
+void common_hal_wifi_radio_stop_station(wifi_radio_obj_t *self) {
+    if (!common_hal_wifi_radio_get_enabled(self)) {
+        mp_raise_RuntimeError(translate("wifi is not enabled"));
+    }
+
+    set_mode_station(self, false);
+}
+
 void common_hal_wifi_radio_start_ap(wifi_radio_obj_t *self, uint8_t *ssid, size_t ssid_len, uint8_t *password, size_t password_len) {
     if (!common_hal_wifi_radio_get_enabled(self)) {
         mp_raise_RuntimeError(translate("wifi is not enabled"));
     }
 
-    start_ap(self);
+    set_mode_ap(self, true);
 
     wifi_config_t *config = &self->ap_config;
     memcpy(&config->ap.ssid, ssid, ssid_len);
@@ -155,10 +185,16 @@ void common_hal_wifi_radio_start_ap(wifi_radio_obj_t *self, uint8_t *ssid, size_
     memcpy(&config->ap.password, password, password_len);
     config->ap.password[password_len] = 0;
     config->ap.authmode = WIFI_AUTH_WPA2_PSK;
-    esp_wifi_set_config(ESP_IF_WIFI_AP, config);
+    config->ap.max_connection = 4;
+    esp_wifi_set_config(WIFI_IF_AP, config);
+}
 
-    // common_hal_wifi_radio_set_enabled(self, false);
-    common_hal_wifi_radio_set_enabled(self, true);
+void common_hal_wifi_radio_stop_ap(wifi_radio_obj_t *self) {
+    if (!common_hal_wifi_radio_get_enabled(self)) {
+        mp_raise_RuntimeError(translate("wifi is not enabled"));
+    }
+
+    set_mode_ap(self, false);
 }
 
 wifi_radio_error_t common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t *ssid, size_t ssid_len, uint8_t *password, size_t password_len, uint8_t channel, mp_float_t timeout, uint8_t *bssid, size_t bssid_len) {
@@ -181,7 +217,7 @@ wifi_radio_error_t common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t
     // explicitly clear bits since xEventGroupWaitBits may have timed out
     xEventGroupClearBits(self->event_group_handle, WIFI_CONNECTED_BIT);
     xEventGroupClearBits(self->event_group_handle, WIFI_DISCONNECTED_BIT);
-    start_station(self);
+    set_mode_station(self, true);
 
     wifi_config_t *config = &self->sta_config;
     memcpy(&config->sta.ssid, ssid, ssid_len);
diff --git a/ports/esp32s2/common-hal/wifi/__init__.c b/ports/esp32s2/common-hal/wifi/__init__.c
index 71cebda9b4..f735d68644 100644
--- a/ports/esp32s2/common-hal/wifi/__init__.c
+++ b/ports/esp32s2/common-hal/wifi/__init__.c
@@ -150,6 +150,9 @@ void common_hal_wifi_init(void) {
     } else if (result != ESP_OK) {
         mp_raise_RuntimeError(translate("Failed to init wifi"));
     }
+    // set station mode to avoid the default SoftAP
+    esp_wifi_set_mode(WIFI_MODE_STA);
+    // start wifi
     common_hal_wifi_radio_set_enabled(self, true);
 }
 
diff --git a/shared-bindings/wifi/Radio.c b/shared-bindings/wifi/Radio.c
index 35cc72fd60..9965455f76 100644
--- a/shared-bindings/wifi/Radio.c
+++ b/shared-bindings/wifi/Radio.c
@@ -170,14 +170,32 @@ const mp_obj_property_t wifi_radio_hostname_obj = {
               (mp_obj_t)&mp_const_none_obj},
 };
 
+//|     def start_station(self) -> None:
+//|     """Starts a Wi-Fi Station"""
+//|     ...
+//|
+STATIC mp_obj_t wifi_radio_start_station(mp_obj_t self) {
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_start_station_obj, wifi_radio_start_station);
+
+//|     def stop_station(self) -> None:
+//|     """Stops the Wi-Fi Station"""
+//|     ...
+//|
+STATIC mp_obj_t wifi_radio_stop_station(mp_obj_t self) {
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_stop_station_obj, wifi_radio_stop_station);
+
 //|     def start_ap(self,
 //|                  ssid: ReadableBuffer,
 //|                  password: ReadableBuffer = b"") -> None:
-//|         """Starts an Access Point with the specified ssid and password."""
+//|         """Starts a Wi-Fi Access Point with the specified ssid and password."""
 //|         ...
 //|
 STATIC mp_obj_t wifi_radio_start_ap(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
-    enum { ARG_ssid, ARG_password, ARG_channel, ARG_bssid, ARG_timeout };
+    enum { ARG_ssid, ARG_password, ARG_channel, ARG_timeout };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_ssid, MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_password,  MP_ARG_OBJ, {.u_obj = MP_OBJ_NULL} },
@@ -204,6 +222,15 @@ STATIC mp_obj_t wifi_radio_start_ap(size_t n_args, const mp_obj_t *pos_args, mp_
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(wifi_radio_start_ap_obj, 1, wifi_radio_start_ap);
 
+//|     def stop_ap(self) -> None:
+//|     """Stops the Wi-Fi Access Point"""
+//|     ...
+//|
+STATIC mp_obj_t wifi_radio_stop_ap(mp_obj_t self) {
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_stop_ap_obj, wifi_radio_stop_ap);
+
 //|     def connect(self,
 //|                 ssid: ReadableBuffer,
 //|                 password: ReadableBuffer = b"",
@@ -417,7 +444,10 @@ STATIC const mp_rom_map_elem_t wifi_radio_locals_dict_table[] = {
 
     { MP_ROM_QSTR(MP_QSTR_hostname),    MP_ROM_PTR(&wifi_radio_hostname_obj) },
 
+    { MP_ROM_QSTR(MP_QSTR_start_station),    MP_ROM_PTR(&wifi_radio_start_station_obj) },
+    { MP_ROM_QSTR(MP_QSTR_stop_station),    MP_ROM_PTR(&wifi_radio_stop_station_obj) },
     { MP_ROM_QSTR(MP_QSTR_start_ap),    MP_ROM_PTR(&wifi_radio_start_ap_obj) },
+    { MP_ROM_QSTR(MP_QSTR_stop_ap),    MP_ROM_PTR(&wifi_radio_stop_ap_obj) },
     { MP_ROM_QSTR(MP_QSTR_connect),    MP_ROM_PTR(&wifi_radio_connect_obj) },
     // { MP_ROM_QSTR(MP_QSTR_connect_to_enterprise),    MP_ROM_PTR(&wifi_radio_connect_to_enterprise_obj) },
 
diff --git a/shared-bindings/wifi/Radio.h b/shared-bindings/wifi/Radio.h
index 61375d4532..b84966f14b 100644
--- a/shared-bindings/wifi/Radio.h
+++ b/shared-bindings/wifi/Radio.h
@@ -83,7 +83,10 @@ extern mp_obj_t common_hal_wifi_radio_get_mac_address_ap(wifi_radio_obj_t *self)
 extern mp_obj_t common_hal_wifi_radio_start_scanning_networks(wifi_radio_obj_t *self);
 extern void common_hal_wifi_radio_stop_scanning_networks(wifi_radio_obj_t *self);
 
+extern void common_hal_wifi_radio_start_station(wifi_radio_obj_t *self);
+extern void common_hal_wifi_radio_stop_station(wifi_radio_obj_t *self);
 extern void common_hal_wifi_radio_start_ap(wifi_radio_obj_t *self, uint8_t *ssid, size_t ssid_len, uint8_t *password, size_t password_len);
+extern void common_hal_wifi_radio_stop_ap(wifi_radio_obj_t *selfn);
 
 extern wifi_radio_error_t common_hal_wifi_radio_connect(wifi_radio_obj_t *self, uint8_t *ssid, size_t ssid_len, uint8_t *password, size_t password_len, uint8_t channel, mp_float_t timeout, uint8_t *bssid, size_t bssid_len);