sta_scan_test: Fix regression, create csv_output option

Signed-off-by: Matthew Stidham <stidmatt@gmail.com>

py-scripts/sta_scan_test.py

@@ -15,6 +15,7 @@ License: Free to distribute and modify. LANforge systems must be licensed.
 import sys
 import os
 import importlib
+import pandas as pd
 
 if sys.version_info[0] != 3:
     print("This script requires Python 3")
@@ -25,8 +26,6 @@ import time
 
 sys.path.append(os.path.join(os.path.abspath(__file__ + "../../../")))
 
-lfcli_base = importlib.import_module("py-json.LANforge.lfcli_base")
-LFCliBase = lfcli_base.LFCliBase
 realm = importlib.import_module("py-json.realm")
 Realm = realm.Realm
 LFUtils = importlib.import_module("py-json.LANforge.LFUtils")
@@ -37,17 +36,20 @@ class StaScan(Realm):
                  ssid=None,
                  security=None,
                  password=None,
-                 sta_list=[],
+                 sta_list=None,
                  upstream=None,
                  radio=None,
                  host="localhost",
                  port=8080,
                  mode=0,
                  number_template="00000",
+                 csv_output=False,
                  use_ht160=False,
                  _debug_on=False,
                  _exit_on_error=False,
                  _exit_on_fail=False):
+        if sta_list is None:
+            sta_list = []
         super().__init__(lfclient_host=host,
                          lfclient_port=port),
         self.upstream = upstream
@@ -60,6 +62,7 @@ class StaScan(Realm):
         self.radio = radio
         self.mode = mode
         self.number_template = number_template
+        self.csv_output = csv_output
         self.debug = _debug_on
         self.station_profile = self.new_station_profile()
         self.station_profile.lfclient_url = self.lfclient_url
@@ -77,22 +80,50 @@ class StaScan(Realm):
     def start(self):
         self.station_profile.admin_up()
         print(self.sta_list)
-        print("Sleeping 15s while waiting for scan")
-        data = {
-            "shelf": 1,
-            "resource": 1,
-            "port": self.sta_list
-        }
-        self.json_post("/cli-json/scan_wifi", data)
-        time.sleep(15)
-        scan_results = self.json_get("scanresults/1/1/%s" % ','.join(self.sta_list))
-
-        print("{0:<23}".format("BSS"), "{0:<7}".format("Signal"), "{0:<5}".format("SSID"))
-        for result in scan_results['scan-results']:
-            for name, info in result.items():
-                print("%s\t%s\t%s" % (info['bss'], info['signal'], info['ssid']))
-
+        LFUtils.wait_until_ports_admin_up(base_url=self.lfclient_url, port_list=self.station_profile.station_names,
+                                          debug_=self.debug)
+        stations = [LFUtils.name_to_eid(x) for x in self.sta_list]
+        stations = pd.DataFrame(stations)
+        resources = stations[1].unique()
+        interfaces = list()
+        for resource in resources:
+            shelf = stations[0][0]
+            resource_station = list(stations[stations[1] == resource][2])
+            url = '/port/%s/%s/%s' % (shelf, resource, ','.join(resource_station))
+            response = self.json_get(url)
+            if 'interface' in response.keys():
+                interface = response['interface']
+                interfaces.append(interface)
+            elif 'interfaces' in response.keys():
+                response_interfaces = response['interfaces']
+                for interface in response_interfaces:
+                    for item in interface.values():
+                        interfaces.append(item)
+        df = pd.DataFrame(interfaces)
+        stations = df[df['port type'] == 'WIFI-STA']
+        stations = list(stations.drop_duplicates('parent dev')['alias'])
+        stations = [station for station in stations if station in self.sta_list]
 
+        for port in stations:
+            port = LFUtils.name_to_eid(port)
+            data = {
+                "shelf": port[0],
+                "resource": port[1],
+                "port": port[2]
+            }
+            self.json_post("/cli-json/scan_wifi", data)
+            time.sleep(15)
+            scan_results = self.json_get("scanresults/%s/%s/%s" % (port[0], port[1], port[2]))
+            if self.csv_output:
+                results = scan_results['scan-results']
+                df = pd.DataFrame([list(result.values())[0] for result in results])
+                df.to_csv(self.csv_output)
+                print('CSV output saved at %s' % self.csv_output)
+            else:
+                print("{0:<23}".format("BSS"), "{0:<7}".format("Signal"), "{0:<5}".format("SSID"))
+                for result in scan_results['scan-results']:
+                    for name, info in result.items():
+                        print("%s\t%s\t%s" % (info['bss'], info['signal'], info['ssid']))
 
     def pre_cleanup(self):
         for sta in self.sta_list:
@@ -133,6 +164,7 @@ def main():
     parser.add_argument('--mode', help='Used to force mode of stations')
     parser.add_argument('--sta_name', help='Optional: User defined station names, can be a comma or space separated list', nargs='+',
                         default=["sta0000"])
+    parser.add_argument('--csv_output', help='create CSV from scan results, otherwise print it in the terminal', default=None)
 
     args = parser.parse_args()
 
@@ -147,6 +179,7 @@ def main():
                             radio=args.radio,
                             security=args.security,
                             use_ht160=False,
+                            csv_output=args.csv_output,
                             mode=args.mode,
                             _debug_on=args.debug)