adafruit · dhalbert · Feb 22, 2024 · Jan 23, 2024 · Jan 23, 2024 · Feb 19, 2024
diff --git a/ports/espressif/common-hal/wifi/Radio.c b/ports/espressif/common-hal/wifi/Radio.c
@@ -250,6 +250,47 @@ void common_hal_wifi_radio_stop_ap(wifi_radio_obj_t *self) {
     set_mode_ap(self, false);
 }
 
+mp_obj_t common_hal_wifi_radio_get_stations_ap(wifi_radio_obj_t *self) {
+    wifi_sta_list_t esp_sta_list;
+    esp_netif_pair_mac_ip_t mac_ip_pair[ESP_WIFI_MAX_CONN_NUM];
+    esp_err_t result;
+
+    result = esp_wifi_ap_get_sta_list(&esp_sta_list);
+    if (result == ESP_OK) {
+        for (int i = 0; i < esp_sta_list.num; i++) {
+            memcpy(mac_ip_pair[i].mac, esp_sta_list.sta[i].mac, MAC_ADDRESS_LENGTH);
+            mac_ip_pair[i].ip.addr = 0;
+        }
+
+        result = esp_netif_dhcps_get_clients_by_mac(self->ap_netif, esp_sta_list.num, mac_ip_pair);
+    }
+
+    if (!self->ap_mode || result != ESP_OK) {
+        return mp_const_none;
+    }
+
+    mp_obj_t mp_sta_list = mp_obj_new_list(0, NULL);
+    for (int i = 0; i < esp_sta_list.num; i++) {
+        mp_obj_t elems[3] = {
+            mp_obj_new_bytes(esp_sta_list.sta[i].mac, MAC_ADDRESS_LENGTH),
+            MP_OBJ_NEW_SMALL_INT(esp_sta_list.sta[i].rssi),
+            mp_const_none
+        };
+
+        if (mac_ip_pair[i].ip.addr) {
+            elems[2] = common_hal_ipaddress_new_ipv4address(mac_ip_pair[i].ip.addr);
+        }
+
+        mp_obj_list_append(mp_sta_list, namedtuple_make_new((const mp_obj_type_t *)&wifi_radio_station_type, 3, 0, elems));
+    }
+
+    return mp_sta_list;
+}
+
+mp_obj_t common_hal_wifi_radio_get_max_stations_ap(wifi_radio_obj_t *self) {
+    return MP_OBJ_NEW_SMALL_INT(ESP_WIFI_MAX_CONN_NUM);
+}
+
 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) {
     if (!common_hal_wifi_radio_get_enabled(self)) {
         mp_raise_RuntimeError(MP_ERROR_TEXT("wifi is not enabled"));

diff --git a/ports/raspberrypi/common-hal/wifi/Radio.c b/ports/raspberrypi/common-hal/wifi/Radio.c
@@ -244,6 +244,59 @@ void common_hal_wifi_radio_stop_ap(wifi_radio_obj_t *self) {
     bindings_cyw43_wifi_enforce_pm();
 }
 
+// There's no published API for the DHCP server to retrieve lease information
+// This code depends on undocumented internal structures and is likely to break in the future
+static uint32_t cyw43_dhcps_get_ip_addr(dhcp_server_t *dhcp_server, uint8_t *mac_address) {
+    for (int i = 0; i < DHCPS_MAX_IP; i++) {
+        if (memcmp(dhcp_server->lease[i].mac, mac_address, MAC_ADDRESS_LENGTH) == 0) {
+            return (dhcp_server->ip.addr & 0x00FFFFFF) + ((DHCPS_BASE_IP + i) << 24);
+        }
+    }
+
+    return 0;
+}
+
+mp_obj_t common_hal_wifi_radio_get_stations_ap(wifi_radio_obj_t *self) {
+    int max_stas;
+    int num_stas;
+
+    if (cyw43_tcpip_link_status(&cyw43_state, CYW43_ITF_AP) != CYW43_LINK_UP) {
+        return mp_const_none;
+    }
+
+    cyw43_wifi_ap_get_max_stas(&cyw43_state, &max_stas);
+
+    uint8_t macs[max_stas * MAC_ADDRESS_LENGTH];
+
+    cyw43_wifi_ap_get_stas(&cyw43_state, &num_stas, macs);
+
+    mp_obj_t mp_sta_list = mp_obj_new_list(0, NULL);
+    for (int i = 0; i < num_stas; i++) {
+        mp_obj_t elems[3] = {
+            mp_obj_new_bytes(&macs[i * MAC_ADDRESS_LENGTH], MAC_ADDRESS_LENGTH),
+            mp_const_none,
+            mp_const_none
+        };
+
+        uint32_t ipv4_addr = cyw43_dhcps_get_ip_addr(&cyw43_state.dhcp_server, &macs[i * MAC_ADDRESS_LENGTH]);
+        if (ipv4_addr) {
+            elems[2] = common_hal_ipaddress_new_ipv4address(ipv4_addr);
+        }
+
+        mp_obj_list_append(mp_sta_list, namedtuple_make_new((const mp_obj_type_t *)&wifi_radio_station_type, 3, 0, elems));
+    }
+
+    return mp_sta_list;
+}
+
+mp_obj_t common_hal_wifi_radio_get_max_stations_ap(wifi_radio_obj_t *self) {
+    int max_stas;
+
+    cyw43_wifi_ap_get_max_stas(&cyw43_state, &max_stas);
+
+    return MP_OBJ_NEW_SMALL_INT(max_stas);
+}
+
 static bool connection_unchanged(wifi_radio_obj_t *self, const uint8_t *ssid, size_t ssid_len) {
     if (cyw43_tcpip_link_status(&cyw43_state, CYW43_ITF_STA) != CYW43_LINK_UP) {
         return false;

diff --git a/shared-bindings/wifi/Radio.c b/shared-bindings/wifi/Radio.c
@@ -661,6 +661,30 @@ STATIC mp_obj_t wifi_radio_get_ap_info(mp_obj_t self) {
 }
 MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_get_ap_info_obj, wifi_radio_get_ap_info);
 
+//|     stations_ap: None
+//|     """In AP mode, returns list of named tuples, each of which contains:
+//|      mac: bytearray (read-only)
+//|      rssi: int (read-only, None on Raspberry Pi Pico W)
+//|      ipv4_address: ipv4_address  (read-only, None if station connected but no address assigned yet or self-assigned address)"""
+STATIC mp_obj_t wifi_radio_get_stations_ap(mp_obj_t self) {
+    return common_hal_wifi_radio_get_stations_ap(self);
+}
+
+MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_get_stations_ap_obj, wifi_radio_get_stations_ap);
+
+MP_PROPERTY_GETTER(wifi_radio_stations_ap_obj,
+    (mp_obj_t)&wifi_radio_get_stations_ap_obj);
+
+//|     max_stations_ap: Optional[Network]
+//|     """In AP mode, returns the maximum possible number of connected stations (read-only)"""
+STATIC mp_obj_t wifi_radio_get_max_stations_ap(mp_obj_t self) {
+    return common_hal_wifi_radio_get_max_stations_ap(self);
+}
+MP_DEFINE_CONST_FUN_OBJ_1(wifi_radio_get_max_stations_ap_obj, wifi_radio_get_max_stations_ap);
+
+MP_PROPERTY_GETTER(wifi_radio_max_stations_ap_obj,
+    (mp_obj_t)&wifi_radio_get_max_stations_ap_obj);
+
 //|     def start_dhcp(self) -> None:
 //|         """Starts the station DHCP client."""
 //|         ...
@@ -750,6 +774,8 @@ STATIC const mp_rom_map_elem_t wifi_radio_locals_dict_table[] = {
     { 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_ap_active),   MP_ROM_PTR(&wifi_radio_ap_active_obj) },
+    { MP_ROM_QSTR(MP_QSTR_stations_ap),   MP_ROM_PTR(&wifi_radio_stations_ap_obj) },
+    { MP_ROM_QSTR(MP_QSTR_max_stations_ap),   MP_ROM_PTR(&wifi_radio_max_stations_ap_obj) },
 
     { MP_ROM_QSTR(MP_QSTR_start_dhcp),    MP_ROM_PTR(&wifi_radio_start_dhcp_client_obj) },
     { MP_ROM_QSTR(MP_QSTR_stop_dhcp),    MP_ROM_PTR(&wifi_radio_stop_dhcp_client_obj) },
@@ -783,3 +809,13 @@ MP_DEFINE_CONST_OBJ_TYPE(
     MP_TYPE_FLAG_HAS_SPECIAL_ACCESSORS,
     locals_dict, &wifi_radio_locals_dict
     );
+
+const mp_obj_namedtuple_type_t wifi_radio_station_type = {
+    NAMEDTUPLE_TYPE_BASE_AND_SLOTS(MP_QSTR_WifiRadioStation),
+    .n_fields = 3,
+    .fields = {
+        MP_QSTR_mac_address,
+        MP_QSTR_rssi,
+        MP_QSTR_ipv4_address,
+    },
+};
diff --git a/shared-bindings/wifi/Radio.h b/shared-bindings/wifi/Radio.h
@@ -32,8 +32,10 @@
 #include "common-hal/wifi/Radio.h"
 
 #include "py/objstr.h"
+#include "py/objnamedtuple.h"
 
 extern const mp_obj_type_t wifi_radio_type;
+extern const mp_obj_namedtuple_type_t wifi_radio_station_type;
 
 typedef enum {
     // 0 is circuitpython-specific; 1-53 are IEEE; 200+ are Espressif
@@ -96,6 +98,8 @@ 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, uint8_t channel, uint32_t authmode, uint8_t max_connections);
 extern void common_hal_wifi_radio_stop_ap(wifi_radio_obj_t *self);
 extern bool common_hal_wifi_radio_get_ap_active(wifi_radio_obj_t *self);
+extern mp_obj_t common_hal_wifi_radio_get_stations_ap(wifi_radio_obj_t *self);
+extern mp_obj_t common_hal_wifi_radio_get_max_stations_ap(wifi_radio_obj_t *self);
 
 extern void common_hal_wifi_radio_start_dhcp_client(wifi_radio_obj_t *self);
 extern void common_hal_wifi_radio_stop_dhcp_client(wifi_radio_obj_t *self);

diff --git a/shared-bindings/wifi/__init__.c b/shared-bindings/wifi/__init__.c
@@ -59,6 +59,7 @@ STATIC const mp_rom_map_elem_t wifi_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_Network),     MP_ROM_PTR(&wifi_network_type) },
     { MP_ROM_QSTR(MP_QSTR_Packet),      MP_ROM_PTR(&wifi_packet_type) },
     { MP_ROM_QSTR(MP_QSTR_Radio),       MP_ROM_PTR(&wifi_radio_type) },
+    { MP_ROM_QSTR(MP_QSTR_Station),     MP_ROM_PTR(&wifi_radio_station_type) },
 
     // Properties
     { MP_ROM_QSTR(MP_QSTR_radio),       MP_ROM_PTR(&common_hal_wifi_radio_obj) },