From 6994d092e1103ec8eaa65214d1d87164856334ac Mon Sep 17 00:00:00 2001
From: Chaitanya Tata <Chaitanya.Tata@nordicsemi.no>
Date: Fri, 26 Jan 2024 21:42:44 +0530
Subject: [PATCH] [nrf fromlist] wifi: shell: Add band support for STA

For a Wi-Fi station the connect API supports both band and channel
configuration, but for a shell command either channel or band makes
sense, so, overload the channel field to support band.

Rejig the band and channel validation to support all modes.

Upstream PR: https://github.com/zephyrproject-rtos/zephyr/pull/68174

Signed-off-by: Chaitanya Tata <Chaitanya.Tata@nordicsemi.no>
---
 subsys/net/l2/wifi/wifi_shell.c | 53 ++++++++++++++++++++++++---------
 1 file changed, 39 insertions(+), 14 deletions(-)

diff --git a/subsys/net/l2/wifi/wifi_shell.c b/subsys/net/l2/wifi/wifi_shell.c
index d70e48024b7..7428a2a79e5 100644
--- a/subsys/net/l2/wifi/wifi_shell.c
+++ b/subsys/net/l2/wifi/wifi_shell.c
@@ -473,7 +473,14 @@ static int __wifi_args_to_params(size_t argc, char *argv[],
 
 	/* Channel (optional: STA, mandatory: AP) */
 	if ((idx < argc) && (strlen(argv[idx]) <= 3)) {
+		uint8_t band;
 		long channel = strtol(argv[idx], &endptr, 10);
+		const uint8_t all_bands[] = {WIFI_FREQ_BAND_2_4_GHZ,
+					WIFI_FREQ_BAND_5_GHZ,
+					WIFI_FREQ_BAND_6_GHZ};
+		bool found = false;
+		char bands_str[MAX_BANDS_STR_LEN] = {0};
+		size_t offset = 0;
 
 		if (*endptr != '\0') {
 			print(context.sh, SHELL_ERROR,
@@ -484,23 +491,40 @@ static int __wifi_args_to_params(size_t argc, char *argv[],
 			return -EINVAL;
 		}
 
-		if (iface_mode == WIFI_MODE_INFRA && channel == 0) {
-			params->channel = WIFI_CHANNEL_ANY;
+		if (iface_mode == WIFI_MODE_INFRA) {
+			if (channel < 0) {
+				/* Negative channel means band */
+				switch (-channel) {
+				case 2:
+					params->band = WIFI_FREQ_BAND_2_4_GHZ;
+					break;
+				case 5:
+					params->band = WIFI_FREQ_BAND_5_GHZ;
+					break;
+				case 6:
+					params->band = WIFI_FREQ_BAND_6_GHZ;
+					break;
+				default:
+					print(context.sh, SHELL_ERROR,
+					      "Invalid band: %ld\n", channel);
+					return -EINVAL;
+				}
+			}
 		} else {
-			const uint8_t bands[] = {WIFI_FREQ_BAND_2_4_GHZ,
-					       WIFI_FREQ_BAND_5_GHZ,
-					       WIFI_FREQ_BAND_6_GHZ};
-			uint8_t band;
-			bool found = false;
-			char bands_str[MAX_BANDS_STR_LEN] = {0};
-			size_t offset = 0;
-
-			for (band = 0; band < ARRAY_SIZE(bands); band++) {
+			if (channel < 0) {
+				print(context.sh, SHELL_ERROR,
+				      "Invalid channel: %ld\n", channel);
+				return -EINVAL;
+			}
+		}
+
+		if (channel > 0) {
+			for (band = 0; band < ARRAY_SIZE(all_bands); band++) {
 				offset += snprintf(bands_str + offset,
 						  sizeof(bands_str) - offset,
 						  "%s%s",
 						  band ? "," : "",
-						  wifi_band_txt(bands[band]));
+						  wifi_band_txt(all_bands[band]));
 				if (offset >= sizeof(bands_str)) {
 					print(context.sh, SHELL_ERROR,
 					      "Failed to parse channel: %s: "
@@ -509,7 +533,7 @@ static int __wifi_args_to_params(size_t argc, char *argv[],
 					return -EINVAL;
 				}
 
-				if (wifi_utils_validate_chan(bands[band],
+				if (wifi_utils_validate_chan(all_bands[band],
 							     channel)) {
 					found = true;
 					break;
@@ -1906,7 +1930,8 @@ SHELL_STATIC_SUBCMD_SET_CREATE(wifi_commands,
 	SHELL_CMD_ARG(connect, NULL,
 		  "Connect to a Wi-Fi AP\n"
 		  "\"<SSID>\"\n"
-		  "[channel number: 0 means all]\n"
+		  "[channel number/band: > 0:Channel, 0:any channel,\n"
+		  "< 0:band (-2:2.4GHz, -5:5GHz, -6:6GHz]\n"
 		  "[PSK: valid only for secure SSIDs]\n"
 		  "[Security type: valid only for secure SSIDs]\n"
 		  "0:None, 1:WPA2-PSK, 2:WPA2-PSK-256, 3:SAE, 4:WAPI, 5:EAP, 6:WEP, 7: WPA-PSK\n"