1# Utils
2# Copyright (c) 2016, Tieto Corporation
3#
4# This software may be distributed under the terms of the BSD license.
5# See README for more details.
6
7import re
8import time
9from remotehost import Host
10import hostapd
11import config
12
13class TestSkip(Exception):
14    def __init__(self, reason):
15        self.reason = reason
16    def __str__(self):
17        return self.reason
18
19# get host based on name
20def get_host(devices, dev_name):
21    dev = config.get_device(devices, dev_name)
22    host = Host(host=dev['hostname'],
23                ifname=dev['ifname'],
24                port=dev['port'],
25                name=dev['name'])
26    host.dev = dev
27    return host
28
29# Run setup_hw - hardware specific
30def setup_hw_host_iface(host, iface, setup_params, force_restart=False):
31    try:
32        setup_hw = setup_params['setup_hw']
33        restart = ""
34        try:
35            if setup_params['restart_device'] == True:
36                restart = "-R"
37        except:
38            pass
39
40        if force_restart:
41            restart = "-R"
42
43        host.execute([setup_hw, "-I", iface, restart])
44    except:
45        pass
46
47def setup_hw_host(host, setup_params, force_restart=False):
48    ifaces = re.split('; | |, ', host.ifname)
49    for iface in ifaces:
50        setup_hw_host_iface(host, iface, setup_params, force_restart)
51
52def setup_hw(hosts, setup_params, force_restart=False):
53    for host in hosts:
54        setup_hw_host(host, setup_params, force_restart)
55
56# get traces - hw specific
57def trace_start(hosts, setup_params):
58    for host in hosts:
59        trace_start_stop(host, setup_params, start=True)
60
61def trace_stop(hosts, setup_params):
62    for host in hosts:
63        trace_start_stop(host, setup_params, start=False)
64
65def trace_start_stop(host, setup_params, start):
66    if setup_params['trace'] == False:
67        return
68    try:
69        start_trace = setup_params['trace_start']
70        stop_trace = setup_params['trace_stop']
71        if start:
72            cmd = start_trace
73        else:
74            cmd = stop_trace
75        trace_dir = setup_params['log_dir'] + host.ifname + "/remote_traces"
76        host.add_log(trace_dir + "/*")
77        host.execute([cmd, "-I", host.ifname, "-D", trace_dir])
78    except:
79        pass
80
81# get perf
82def perf_start(hosts, setup_params):
83    for host in hosts:
84        perf_start_stop(host, setup_params, start=True)
85
86def perf_stop(hosts, setup_params):
87    for host in hosts:
88        perf_start_stop(host, setup_params, start=False)
89
90def perf_start_stop(host, setup_params, start):
91    if setup_params['perf'] == False:
92        return
93    try:
94        perf_start = setup_params['perf_start']
95        perf_stop = setup_params['perf_stop']
96        if start:
97            cmd = perf_start
98        else:
99            cmd = perf_stop
100        perf_dir = setup_params['log_dir'] + host.ifname + "/remote_perf"
101        host.add_log(perf_dir + "/*")
102        host.execute([cmd, "-I", host.ifname, "-D", perf_dir])
103    except:
104        pass
105
106# hostapd/wpa_supplicant helpers
107def run_hostapd(host, setup_params):
108    log_file = None
109    try:
110        tc_name = setup_params['tc_name']
111        log_dir = setup_params['log_dir']
112        log_file = log_dir + tc_name + "_hostapd_" + host.name + "_" + host.ifname + ".log"
113        host.execute(["rm", log_file])
114        log = " -f " + log_file
115    except:
116        log = ""
117
118    if log_file:
119        host.add_log(log_file)
120    pidfile = setup_params['log_dir'] + "hostapd_" + host.ifname + "_" + setup_params['tc_name'] + ".pid"
121    status, buf = host.execute([setup_params['hostapd'], "-B", "-ddt", "-g", "udp:" + host.port, "-P", pidfile, log])
122    if status != 0:
123        raise Exception("Could not run hostapd: " + buf)
124
125def run_wpasupplicant(host, setup_params):
126    log_file = None
127    try:
128        tc_name = setup_params['tc_name']
129        log_dir = setup_params['log_dir']
130        log_file = log_dir + tc_name + "_wpa_supplicant_" + host.name + "_" + host.ifname + ".log"
131        host.execute(["rm", log_file])
132        log = " -f " + log_file
133    except:
134        log = ""
135
136    if log_file:
137        host.add_log(log_file)
138    pidfile = setup_params['log_dir'] + "wpa_supplicant_" + host.ifname + "_" + setup_params['tc_name'] + ".pid"
139    status, buf = host.execute([setup_params['wpa_supplicant'], "-B", "-ddt", "-g", "udp:" + host.port, "-P", pidfile, log])
140    if status != 0:
141        raise Exception("Could not run wpa_supplicant: " + buf)
142
143def kill_wpasupplicant(host, setup_params):
144    pidfile = setup_params['log_dir'] + "wpa_supplicant_" + host.ifname + "_" + setup_params['tc_name'] + ".pid"
145    host.execute(["kill `cat " + pidfile + "`"])
146
147def kill_hostapd(host, setup_params):
148    pidfile = setup_params['log_dir'] + "hostapd_" + host.ifname + "_" + setup_params['tc_name'] + ".pid"
149    host.execute(["kill `cat " + pidfile + "`"])
150
151def get_ap_params(channel="1", bw="HT20", country="US", security="open", ht_capab=None, vht_capab=None):
152    ssid = "test_" + channel + "_" + security + "_" + bw
153
154    if bw == "b_only":
155        params = hostapd.b_only_params(channel, ssid, country)
156    elif bw == "g_only":
157        params = hostapd.g_only_params(channel, ssid, country)
158    elif bw == "g_only_wmm":
159        params = hostapd.g_only_params(channel, ssid, country)
160        params['wmm_enabled'] = "1"
161    elif bw == "a_only":
162        params = hostapd.a_only_params(channel, ssid, country)
163    elif bw == "a_only_wmm":
164        params = hostapd.a_only_params(channel, ssid, country)
165        params['wmm_enabled'] = "1"
166    elif bw == "HT20":
167        params = hostapd.ht20_params(channel, ssid, country)
168        if ht_capab:
169            try:
170                params['ht_capab'] = params['ht_capab'] + ht_capab
171            except:
172                params['ht_capab'] = ht_capab
173    elif bw == "HT40+":
174        params = hostapd.ht40_plus_params(channel, ssid, country)
175        if ht_capab:
176            params['ht_capab'] = params['ht_capab'] + ht_capab
177    elif bw == "HT40-":
178        params = hostapd.ht40_minus_params(channel, ssid, country)
179        if ht_capab:
180            params['ht_capab'] = params['ht_capab'] + ht_capab
181    elif bw == "VHT80":
182        params = hostapd.ht40_plus_params(channel, ssid, country)
183        if ht_capab:
184            params['ht_capab'] = params['ht_capab'] + ht_capab
185        if vht_capab:
186            try:
187                params['vht_capab'] = params['vht_capab'] + vht_capab
188            except:
189                params['vht_capab'] = vht_capab
190        params['ieee80211ac'] = "1"
191        params['vht_oper_chwidth'] = "1"
192        params['vht_oper_centr_freq_seg0_idx'] = str(int(channel) + 6)
193    else:
194        params = {}
195
196    # now setup security params
197    if security == "tkip":
198        sec_params = hostapd.wpa_params(passphrase="testtest")
199    elif security == "ccmp":
200        sec_params = hostapd.wpa2_params(passphrase="testtest")
201    elif security == "mixed":
202        sec_params = hostapd.wpa_mixed_params(passphrase="testtest")
203    elif security == "wep":
204        sec_params = {"wep_key0" : "123456789a",
205                      "wep_default_key" : "0",
206                      "auth_algs" : "1"}
207    elif security == "wep_shared":
208        sec_params = {"wep_key0" : "123456789a",
209                      "wep_default_key" : "0",
210                      "auth_algs" : "2"}
211    else:
212        sec_params = {}
213
214    params.update(sec_params)
215
216    return params
217
218# ip helpers
219def get_ipv4(client, ifname=None):
220    if ifname is None:
221        ifname = client.ifname
222    status, buf = client.execute(["ifconfig", ifname])
223    lines = buf.splitlines()
224
225    for line in lines:
226        res = line.find("inet addr:")
227        if res != -1:
228            break
229
230    if res != -1:
231        words = line.split()
232        addr = words[1].split(":")
233        return addr[1]
234
235    return "unknown"
236
237def get_ipv6(client, ifname=None):
238    res = -1
239    if ifname is None:
240        ifname = client.ifname
241    status, buf = client.execute(["ifconfig", ifname])
242    lines = buf.splitlines()
243
244    for line in lines:
245        res = line.find("Scope:Link")
246        if res == -1:
247            res = line.find("<link>")
248        if res != -1:
249            break
250
251    if res != -1:
252        words = line.split()
253        if words[0] == "inet6" and words[1] == "addr:":
254            addr_mask = words[2]
255            addr = addr_mask.split("/")
256            return addr[0]
257        if words[0] == "inet6":
258            return words[1]
259
260    return "unknown"
261
262def get_ip(client, addr_type="ipv6", iface=None):
263    if addr_type == "ipv6":
264        return get_ipv6(client, iface)
265    elif addr_type == "ipv4":
266        return get_ipv4(client, iface)
267    else:
268        return "unknown addr_type: " + addr_type
269
270def get_ipv4_addr(setup_params, number):
271    try:
272        ipv4_base = setup_params['ipv4_test_net']
273    except:
274        ipv4_base = "172.16.12.0"
275
276    parts = ipv4_base.split('.')
277    ipv4 = parts[0] + "." + parts[1] + "." + parts[2] + "." + str(number)
278
279    return ipv4
280
281def get_mac_addr(host, iface=None):
282    if iface == None:
283        iface = host.ifname
284    status, buf = host.execute(["ifconfig", iface])
285    if status != 0:
286        raise Exception("ifconfig " + iface)
287    words = buf.split()
288    found = 0
289    for word in words:
290        if found == 1:
291            return word
292        if word == "HWaddr" or word == "ether":
293            found = 1
294    raise Exception("Could not find HWaddr")
295
296# connectivity/ping helpers
297def get_ping_packet_loss(ping_res):
298    loss_line = ""
299    lines = ping_res.splitlines()
300    for line in lines:
301        if line.find("packet loss") != -1:
302            loss_line = line
303            break;
304
305    if loss_line == "":
306        return "100%"
307
308    sections = loss_line.split(",")
309
310    for section in sections:
311        if section.find("packet loss") != -1:
312            words = section.split()
313            return words[0]
314
315    return "100%"
316
317def ac_to_ping_ac(qos):
318    if qos == "be":
319        qos_param = "0x00"
320    elif qos == "bk":
321        qos_param = "0x20"
322    elif qos == "vi":
323        qos_param = "0xA0"
324    elif qos == "vo":
325        qos_param = "0xE0"
326    else:
327        qos_param = "0x00"
328    return qos_param
329
330def ping_run(host, ip, result, ifname=None, addr_type="ipv4", deadline="5", qos=None):
331    if ifname is None:
332       ifname = host.ifname
333    if addr_type == "ipv6":
334        ping = ["ping6"]
335    else:
336        ping = ["ping"]
337
338    ping = ping + ["-w", deadline, "-I", ifname]
339    if qos:
340        ping = ping + ["-Q", ac_to_ping_ac(qos)]
341    ping = ping + [ip]
342
343    flush_arp_cache(host)
344
345    thread = host.thread_run(ping, result)
346    return thread
347
348def ping_wait(host, thread, timeout=None):
349    host.thread_wait(thread, timeout)
350    if thread.is_alive():
351        raise Exception("ping thread still alive")
352
353def flush_arp_cache(host):
354    host.execute(["ip", "-s", "-s", "neigh", "flush", "all"])
355
356def check_connectivity(a, b, addr_type="ipv4", deadline="5", qos=None):
357    addr_a = get_ip(a, addr_type)
358    addr_b = get_ip(b, addr_type)
359
360    if addr_type == "ipv4":
361        ping = ["ping"]
362    else:
363        ping = ["ping6"]
364
365    ping_a_b = ping + ["-w", deadline, "-I", a.ifname]
366    ping_b_a = ping + ["-w", deadline, "-I", b.ifname]
367    if qos:
368        ping_a_b = ping_a_b + ["-Q", ac_to_ping_ac(qos)]
369        ping_b_a = ping_b_a + ["-Q", ac_to_ping_ac(qos)]
370    ping_a_b = ping_a_b + [addr_b]
371    ping_b_a = ping_b_a + [addr_a]
372
373    # Clear arp cache
374    flush_arp_cache(a)
375    flush_arp_cache(b)
376
377    status, buf = a.execute(ping_a_b)
378    if status == 2 and ping == "ping6":
379        # tentative possible for a while, try again
380        time.sleep(3)
381        status, buf = a.execute(ping_a_b)
382    if status != 0:
383        raise Exception("ping " + a.name + "/" + a.ifname + " >> " + b.name + "/" + b.ifname)
384
385    a_b = get_ping_packet_loss(buf)
386
387    # Clear arp cache
388    flush_arp_cache(a)
389    flush_arp_cache(b)
390
391    status, buf = b.execute(ping_b_a)
392    if status != 0:
393        raise Exception("ping " + b.name + "/" + b.ifname + " >> " + a.name + "/" + a.ifname)
394
395    b_a = get_ping_packet_loss(buf)
396
397    if int(a_b[:-1]) > 40:
398        raise Exception("Too high packet lost: " + a_b)
399
400    if int(b_a[:-1]) > 40:
401        raise Exception("Too high packet lost: " + b_a)
402
403    return a_b, b_a
404
405
406# iperf helpers
407def get_iperf_speed(iperf_res, pattern="Mbits/sec"):
408    lines = iperf_res.splitlines()
409    sum_line = ""
410    last_line = ""
411    count = 0
412    res = -1
413
414    # first find last SUM line
415    for line in lines:
416        res = line.find("[SUM]")
417        if res != -1:
418            sum_line = line
419
420    # next check SUM status
421    if sum_line != "":
422        words = sum_line.split()
423        for word in words:
424            res = word.find(pattern)
425            if res != -1:
426                return words[count - 1] + " " + pattern
427            count = count + 1
428
429    # no SUM - one thread - find last line
430    for line in lines:
431        res = line.find(pattern)
432        if res != -1:
433            last_line = line
434
435    if last_line == "":
436        return "0 " + pattern
437
438    count = 0
439    words = last_line.split()
440    for word in words:
441        res = word.find(pattern)
442        if res != -1:
443            return words[count - 1] + " " + pattern
444            break;
445        count = count + 1
446    return "0 " + pattern
447
448def ac_to_iperf_ac(qos):
449    if qos == "be":
450        qos_param = "0x00"
451    elif qos == "bk":
452        qos_param = "0x20"
453    elif qos == "vi":
454        qos_param = "0xA0"
455    elif qos == "vo":
456        qos_param = "0xE0"
457    else:
458        qos_param = "0x00"
459    return qos_param
460
461def iperf_run(server, client, server_ip, client_res, server_res,
462              l4="udp", bw="30M", test_time="30", parallel="5",
463              qos="be", param=" -i 5 ", ifname=None, l3="ipv4",
464              port="5001", iperf="iperf"):
465    if ifname == None:
466        ifname = client.ifname
467
468    if iperf == "iperf":
469        iperf_server = [iperf]
470    elif iperf == "iperf3":
471        iperf_server = [iperf, "-1"]
472
473    if l3 == "ipv4":
474        iperf_client = [iperf, "-c", server_ip, "-p", port]
475        iperf_server = iperf_server + ["-p", port]
476    elif l3 == "ipv6":
477        iperf_client = [iperf, "-V", "-c", server_ip  + "%" + ifname, "-p", port]
478        iperf_server = iperf_server + ["-V", "-p", port]
479    else:
480        return -1, -1
481
482    iperf_server = iperf_server + ["-s", "-f", "m", param]
483    iperf_client = iperf_client + ["-f", "m", "-t", test_time]
484
485    if parallel != "1":
486        iperf_client = iperf_client + ["-P", parallel]
487
488    if l4 == "udp":
489        if iperf != "iperf3":
490            iperf_server = iperf_server + ["-u"]
491        iperf_client = iperf_client + ["-u", "-b", bw]
492
493    if qos:
494        iperf_client = iperf_client + ["-Q", ac_to_iperf_ac(qos)]
495
496    flush_arp_cache(server)
497    flush_arp_cache(client)
498
499    server_thread = server.thread_run(iperf_server, server_res)
500    time.sleep(1)
501    client_thread = client.thread_run(iperf_client, client_res)
502
503    return server_thread, client_thread
504
505def iperf_wait(server, client, server_thread, client_thread, timeout=None, iperf="iperf"):
506    client.thread_wait(client_thread, timeout)
507    if client_thread.is_alive():
508        raise Exception("iperf client thread still alive")
509
510    server.thread_wait(server_thread, 5)
511    if server_thread.is_alive():
512        server.execute(["killall", "-s", "INT", iperf])
513        time.sleep(1)
514
515    server.thread_wait(server_thread, 5)
516    if server_thread.is_alive():
517        raise Exception("iperf server thread still alive")
518
519    return
520
521def run_tp_test(server, client, l3="ipv4", iperf="iperf", l4="tcp", test_time="10", parallel="5",
522                qos="be", bw="30M", ifname=None, port="5001"):
523    client_res = []
524    server_res = []
525
526    server_ip = get_ip(server, l3)
527    time.sleep(1)
528    server_thread, client_thread = iperf_run(server, client, server_ip, client_res, server_res,
529                                             l3=l3, iperf=iperf, l4=l4, test_time=test_time,
530                                             parallel=parallel, qos=qos, bw=bw, ifname=ifname,
531                                             port=port)
532    iperf_wait(server, client, server_thread, client_thread, iperf=iperf, timeout=int(test_time) + 10)
533
534    if client_res[0] != 0:
535        raise Exception(iperf + " client: " + client_res[1])
536    if server_res[0] != 0:
537        raise Exception(iperf + " server: " + server_res[1])
538    if client_res[1] is None:
539        raise Exception(iperf + " client result issue")
540    if server_res[1] is None:
541        raise Exception(iperf + " server result issue")
542
543    if iperf == "iperf":
544          result = server_res[1]
545    if iperf == "iperf3":
546          result = client_res[1]
547
548    speed = get_iperf_speed(result)
549    return speed
550
551def get_iperf_bw(bw, parallel, spacial_streams=2):
552    if bw == "b_only":
553        max_tp = 11
554    elif bw == "g_only" or bw == "g_only_wmm" or bw == "a_only" or bw == "a_only_wmm":
555        max_tp = 54
556    elif bw == "HT20":
557        max_tp = 72 * spacial_streams
558    elif bw == "HT40+" or bw == "HT40-":
559        max_tp = 150 * spacial_streams
560    elif bw == "VHT80":
561        max_tp = 433 * spacial_streams
562    else:
563        max_tp = 150
564
565    max_tp = 1.2 * max_tp
566
567    return str(int(max_tp/int(parallel))) + "M"
568