trex_console.py revision 87bac1ab
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4"""
5Dan Klein, Itay Marom
6Cisco Systems, Inc.
7
8Copyright (c) 2015-2015 Cisco Systems, Inc.
9Licensed under the Apache License, Version 2.0 (the "License");
10you may not use this file except in compliance with the License.
11You may obtain a copy of the License at
12    http://www.apache.org/licenses/LICENSE-2.0
13Unless required by applicable law or agreed to in writing, software
14distributed under the License is distributed on an "AS IS" BASIS,
15WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16See the License for the specific language governing permissions and
17limitations under the License.
18"""
19from __future__ import print_function
20
21import subprocess
22import cmd
23import json
24import ast
25import argparse
26import random
27import readline
28import string
29import os
30import sys
31import tty, termios
32
33try:
34    import stl_path
35except:
36    from . import stl_path
37from trex_stl_lib.api import *
38
39from trex_stl_lib.utils.text_opts import *
40from trex_stl_lib.utils.common import user_input, get_current_user
41from trex_stl_lib.utils import parsing_opts
42
43try:
44    import trex_tui
45except:
46    from . import trex_tui
47
48from functools import wraps
49
50__version__ = "1.1"
51
52# console custom logger
53class ConsoleLogger(LoggerApi):
54    def __init__ (self):
55        self.prompt_redraw = None
56
57    def write (self, msg, newline = True):
58        if newline:
59            print(msg)
60        else:
61            print(msg, end=' ')
62
63    def flush (self):
64        sys.stdout.flush()
65
66    # override this for the prompt fix
67    def async_log (self, msg, level = LoggerApi.VERBOSE_REGULAR, newline = True):
68        self.log(msg, level, newline)
69        if ( (self.level >= LoggerApi.VERBOSE_REGULAR) and self.prompt_redraw ):
70            self.prompt_redraw()
71            self.flush()
72
73
74def set_window_always_on_top (title):
75    # we need the GDK module, if not available - ignroe this command
76    try:
77        if sys.version_info < (3,0):
78            from gtk import gdk
79        else:
80            #from gi.repository import Gdk as gdk
81            return
82
83    except ImportError:
84        return
85
86    # search the window and set it as above
87    root = gdk.get_default_root_window()
88
89    for id in root.property_get('_NET_CLIENT_LIST')[2]:
90        w = gdk.window_foreign_new(id)
91        if w:
92            name = w.property_get('WM_NAME')[2]
93            if name == title:
94                w.set_keep_above(True)
95                gdk.window_process_all_updates()
96                break
97
98
99class TRexGeneralCmd(cmd.Cmd):
100    def __init__(self):
101        cmd.Cmd.__init__(self)
102        # configure history behaviour
103        self._history_file_dir = "/tmp/trex/console/"
104        self._history_file = self.get_history_file_full_path()
105        readline.set_history_length(100)
106        # load history, if any
107        self.load_console_history()
108
109
110    def get_console_identifier(self):
111        return self.__class__.__name__
112
113    def get_history_file_full_path(self):
114        return "{dir}{filename}.hist".format(dir=self._history_file_dir,
115                                             filename=self.get_console_identifier())
116
117    def load_console_history(self):
118        if os.path.exists(self._history_file):
119            readline.read_history_file(self._history_file)
120        return
121
122    def save_console_history(self):
123        if not os.path.exists(self._history_file_dir):
124            # make the directory available for every user
125            try:
126                original_umask = os.umask(0)
127                os.makedirs(self._history_file_dir, mode = 0o777)
128            finally:
129                os.umask(original_umask)
130
131
132        # os.mknod(self._history_file)
133        readline.write_history_file(self._history_file)
134        return
135
136    def print_history (self):
137
138        length = readline.get_current_history_length()
139
140        for i in range(1, length + 1):
141            cmd = readline.get_history_item(i)
142            print("{:<5}   {:}".format(i, cmd))
143
144    def get_history_item (self, index):
145        length = readline.get_current_history_length()
146        if index > length:
147            print(format_text("please select an index between {0} and {1}".format(0, length)))
148            return None
149
150        return readline.get_history_item(index)
151
152
153    def emptyline(self):
154        """Called when an empty line is entered in response to the prompt.
155
156        This overriding is such that when empty line is passed, **nothing happens**.
157        """
158        return
159
160    def completenames(self, text, *ignored):
161        """
162        This overriding is such that a space is added to name completion.
163        """
164        dotext = 'do_'+text
165        return [a[3:]+' ' for a in self.get_names() if a.startswith(dotext)]
166
167
168#
169# main console object
170class TRexConsole(TRexGeneralCmd):
171    """Trex Console"""
172
173    def __init__(self, stateless_client, verbose = False):
174
175        self.stateless_client = stateless_client
176
177        TRexGeneralCmd.__init__(self)
178
179        self.tui = trex_tui.TrexTUI(stateless_client)
180        self.terminal = None
181
182        self.verbose = verbose
183
184        self.intro  = "\n-=TRex Console v{ver}=-\n".format(ver=__version__)
185        self.intro += "\nType 'help' or '?' for supported actions\n"
186
187        self.postcmd(False, "")
188
189
190    ################### internal section ########################
191
192    def prompt_redraw (self):
193        self.postcmd(False, "")
194        sys.stdout.write("\n" + self.prompt + readline.get_line_buffer())
195        sys.stdout.flush()
196
197
198    def verify_connected(f):
199        @wraps(f)
200        def wrap(*args):
201            inst = args[0]
202            func_name = f.__name__
203            if func_name.startswith("do_"):
204                func_name = func_name[3:]
205
206            if not inst.stateless_client.is_connected():
207                print(format_text("\n'{0}' cannot be executed on offline mode\n".format(func_name), 'bold'))
208                return
209
210            ret = f(*args)
211            return ret
212
213        return wrap
214
215    # TODO: remove this ugly duplication
216    def verify_connected_and_rw (f):
217        @wraps(f)
218        def wrap(*args):
219            inst = args[0]
220            func_name = f.__name__
221            if func_name.startswith("do_"):
222                func_name = func_name[3:]
223
224            if not inst.stateless_client.is_connected():
225                print(format_text("\n'{0}' cannot be executed on offline mode\n".format(func_name), 'bold'))
226                return
227
228            if inst.stateless_client.is_all_ports_acquired():
229                print(format_text("\n'{0}' cannot be executed on read only mode\n".format(func_name), 'bold'))
230                return
231
232            rc = f(*args)
233            return rc
234
235        return wrap
236
237
238    def get_console_identifier(self):
239        return "{context}_{server}".format(context=get_current_user(),
240                                           server=self.stateless_client.get_connection_info()['server'])
241
242    def register_main_console_methods(self):
243        main_names = set(self.trex_console.get_names()).difference(set(dir(self.__class__)))
244        for name in main_names:
245            for prefix in 'do_', 'help_', 'complete_':
246                if name.startswith(prefix):
247                    self.__dict__[name] = getattr(self.trex_console, name)
248
249    def precmd(self, line):
250        # before doing anything, save history snapshot of the console
251        # this is done before executing the command in case of ungraceful application exit
252        self.save_console_history()
253
254        lines = line.split(';')
255
256        for line in lines:
257            stop = self.onecmd(line)
258            stop = self.postcmd(stop, line)
259            if stop:
260                return "quit"
261
262        return ""
263
264
265    def postcmd(self, stop, line):
266
267        if not self.stateless_client.is_connected():
268            self.prompt = "trex(offline)>"
269            self.supported_rpc = None
270            return stop
271
272        if self.stateless_client.is_all_ports_acquired():
273            self.prompt = "trex(read-only)>"
274            return stop
275
276
277        self.prompt = "trex>"
278
279        return stop
280
281    def default(self, line):
282        print("'{0}' is an unrecognized command. type 'help' or '?' for a list\n".format(line))
283
284    @staticmethod
285    def tree_autocomplete(text):
286        dir = os.path.dirname(text)
287        if dir:
288            path = dir
289        else:
290            path = "."
291
292
293        start_string = os.path.basename(text)
294
295        targets = []
296
297        for x in os.listdir(path):
298            if x.startswith(start_string):
299                y = os.path.join(path, x)
300                if os.path.isfile(y):
301                    targets.append(x + ' ')
302                elif os.path.isdir(y):
303                    targets.append(x + '/')
304
305        return targets
306
307
308    ####################### shell commands #######################
309    @verify_connected
310    def do_ping (self, line):
311        '''Ping the server\n'''
312        self.stateless_client.ping()
313
314
315    # set verbose on / off
316    def do_verbose(self, line):
317        '''Shows or set verbose mode\n'''
318        if line == "":
319            print("\nverbose is " + ("on\n" if self.verbose else "off\n"))
320
321        elif line == "on":
322            self.verbose = True
323            self.stateless_client.set_verbose("high")
324            print(format_text("\nverbose set to on\n", 'green', 'bold'))
325
326        elif line == "off":
327            self.verbose = False
328            self.stateless_client.set_verbose("normal")
329            print(format_text("\nverbose set to off\n", 'green', 'bold'))
330
331        else:
332            print(format_text("\nplease specify 'on' or 'off'\n", 'bold'))
333
334    # show history
335    def help_history (self):
336        self.do_history("-h")
337
338    def do_shell (self, line):
339        return self.do_history(line)
340
341    def do_push (self, line):
342        '''Push a PCAP file\n'''
343        return self.stateless_client.push_line(line)
344
345    def help_push (self):
346        return self.do_push("-h")
347
348    def do_portattr (self, line):
349        '''Change/show port(s) attributes\n'''
350        return self.stateless_client.set_port_attr_line(line)
351
352    def help_portattr (self):
353        return self.do_portattr("-h")
354
355    @verify_connected
356    def do_map (self, line):
357        '''Maps ports topology\n'''
358        ports = self.stateless_client.get_acquired_ports()
359        if not ports:
360            print("No ports acquired\n")
361            return
362
363        with self.stateless_client.logger.supress():
364            table = stl_map_ports(self.stateless_client, ports = ports)
365
366
367        print(format_text('\nAcquired ports topology:\n', 'bold', 'underline'))
368
369        # bi-dir ports
370        print(format_text('Bi-directional ports:\n','underline'))
371        for port_a, port_b in table['bi']:
372            print("port {0} <--> port {1}".format(port_a, port_b))
373
374        print("")
375
376        # unknown ports
377        print(format_text('Mapping unknown:\n','underline'))
378        for port in table['unknown']:
379            print("port {0}".format(port))
380        print("")
381
382
383
384
385    def do_history (self, line):
386        '''Manage the command history\n'''
387
388        item = parsing_opts.ArgumentPack(['item'],
389                                         {"nargs": '?',
390                                          'metavar': 'item',
391                                          'type': parsing_opts.check_negative,
392                                          'help': "an history item index",
393                                          'default': 0})
394
395        parser = parsing_opts.gen_parser(self,
396                                         "history",
397                                         self.do_history.__doc__,
398                                         item)
399
400        opts = parser.parse_args(line.split())
401        if opts is None:
402            return
403
404        if opts.item == 0:
405            self.print_history()
406        else:
407            cmd = self.get_history_item(opts.item)
408            if cmd == None:
409                return
410
411            print("Executing '{0}'".format(cmd))
412
413            return self.onecmd(cmd)
414
415
416
417    ############### connect
418    def do_connect (self, line):
419        '''Connects to the server\n'''
420
421        self.stateless_client.connect_line(line)
422
423
424    def do_disconnect (self, line):
425        '''Disconnect from the server\n'''
426
427        self.stateless_client.disconnect_line(line)
428
429
430    ############### start
431
432    def complete_start(self, text, line, begidx, endidx):
433        s = line.split()
434        l = len(s)
435
436        file_flags = parsing_opts.get_flags(parsing_opts.FILE_PATH)
437
438        if (l > 1) and (s[l - 1] in file_flags):
439            return TRexConsole.tree_autocomplete("")
440
441        if (l > 2) and (s[l - 2] in file_flags):
442            return TRexConsole.tree_autocomplete(s[l - 1])
443
444
445    @verify_connected_and_rw
446    def do_start(self, line):
447        '''Start selected traffic in specified port(s) on TRex\n'''
448
449        self.stateless_client.start_line(line)
450
451
452
453
454    def help_start(self):
455        self.do_start("-h")
456
457    ############# stop
458    @verify_connected_and_rw
459    def do_stop(self, line):
460        '''stops port(s) transmitting traffic\n'''
461
462        self.stateless_client.stop_line(line)
463
464    def help_stop(self):
465        self.do_stop("-h")
466
467    ############# update
468    @verify_connected_and_rw
469    def do_update(self, line):
470        '''update speed of port(s)currently transmitting traffic\n'''
471
472        self.stateless_client.update_line(line)
473
474    def help_update (self):
475        self.do_update("-h")
476
477    ############# pause
478    @verify_connected_and_rw
479    def do_pause(self, line):
480        '''pause port(s) transmitting traffic\n'''
481
482        self.stateless_client.pause_line(line)
483
484    ############# resume
485    @verify_connected_and_rw
486    def do_resume(self, line):
487        '''resume port(s) transmitting traffic\n'''
488
489        self.stateless_client.resume_line(line)
490
491
492
493    ########## reset
494    @verify_connected_and_rw
495    def do_reset (self, line):
496        '''force stop all ports\n'''
497        self.stateless_client.reset_line(line)
498
499
500    ######### validate
501    @verify_connected
502    def do_validate (self, line):
503        '''validates port(s) stream configuration\n'''
504
505        self.stateless_client.validate_line(line)
506
507
508    @verify_connected
509    def do_stats(self, line):
510        '''Fetch statistics from TRex server by port\n'''
511        self.stateless_client.show_stats_line(line)
512
513
514    def help_stats(self):
515        self.do_stats("-h")
516
517    @verify_connected
518    def do_streams(self, line):
519        '''Fetch statistics from TRex server by port\n'''
520        self.stateless_client.show_streams_line(line)
521
522
523    def help_streams(self):
524        self.do_streams("-h")
525
526    @verify_connected
527    def do_clear(self, line):
528        '''Clear cached local statistics\n'''
529        self.stateless_client.clear_stats_line(line)
530
531
532    def help_clear(self):
533        self.do_clear("-h")
534
535
536    def help_events (self):
537        self.do_events("-h")
538
539    def do_events (self, line):
540        '''shows events recieved from server\n'''
541        return self.stateless_client.get_events_line(line)
542
543
544    def complete_profile(self, text, line, begidx, endidx):
545        return self.complete_start(text,line, begidx, endidx)
546
547    def do_profile (self, line):
548        '''shows information about a profile'''
549        self.stateless_client.show_profile_line(line)
550
551    # tui
552    @verify_connected
553    def do_tui (self, line):
554        '''Shows a graphical console\n'''
555
556        parser = parsing_opts.gen_parser(self,
557                                         "tui",
558                                         self.do_tui.__doc__,
559                                         parsing_opts.XTERM)
560
561        opts = parser.parse_args(line.split())
562        if opts is None:
563            return
564
565        if opts.xterm:
566            if not os.path.exists('/usr/bin/xterm'):
567                print(format_text("XTERM does not exists on this machine", 'bold'))
568                return
569
570            info = self.stateless_client.get_connection_info()
571
572            exe = './trex-console --top -t -q -s {0} -p {1} --async_port {2}'.format(info['server'], info['sync_port'], info['async_port'])
573            cmd = ['/usr/bin/xterm', '-geometry', '111x48', '-sl', '0', '-title', 'trex_tui', '-e', exe]
574
575            # detach child
576            self.terminal = subprocess.Popen(cmd, preexec_fn = os.setpgrp)
577
578            return
579
580
581        with self.stateless_client.logger.supress():
582            self.tui.show()
583
584
585    def help_tui (self):
586        do_tui("-h")
587
588
589    # quit function
590    def do_quit(self, line):
591        '''Exit the client\n'''
592        return True
593
594
595    def do_help (self, line):
596         '''Shows This Help Screen\n'''
597         if line:
598             try:
599                 func = getattr(self, 'help_' + line)
600             except AttributeError:
601                 try:
602                     doc = getattr(self, 'do_' + line).__doc__
603                     if doc:
604                         self.stdout.write("%s\n"%str(doc))
605                         return
606                 except AttributeError:
607                     pass
608                 self.stdout.write("%s\n"%str(self.nohelp % (line,)))
609                 return
610             func()
611             return
612
613         print("\nSupported Console Commands:")
614         print("----------------------------\n")
615
616         cmds =  [x[3:] for x in self.get_names() if x.startswith("do_")]
617         hidden = ['EOF', 'q', 'exit', 'h', 'shell']
618         for cmd in cmds:
619             if cmd in hidden:
620                 continue
621
622             try:
623                 doc = getattr(self, 'do_' + cmd).__doc__
624                 if doc:
625                     help = str(doc)
626                 else:
627                     help = "*** Undocumented Function ***\n"
628             except AttributeError:
629                 help = "*** Undocumented Function ***\n"
630
631             l=help.splitlines()
632             print("{:<30} {:<30}".format(cmd + " - ",l[0] ))
633
634    # a custorm cmdloop wrapper
635    def start(self):
636        while True:
637            try:
638                self.cmdloop()
639                break
640            except KeyboardInterrupt as e:
641                if not readline.get_line_buffer():
642                    raise KeyboardInterrupt
643                else:
644                    print("")
645                    self.intro = None
646                    continue
647
648        if self.terminal:
649            self.terminal.kill()
650
651    # aliases
652    do_exit = do_EOF = do_q = do_quit
653    do_h = do_history
654
655
656# run a script of commands
657def run_script_file (self, filename, stateless_client):
658
659    self.logger.log(format_text("\nRunning script file '{0}'...".format(filename), 'bold'))
660
661    with open(filename) as f:
662        script_lines = f.readlines()
663
664    cmd_table = {}
665
666    # register all the commands
667    cmd_table['start'] = stateless_client.start_line
668    cmd_table['stop']  = stateless_client.stop_line
669    cmd_table['reset'] = stateless_client.reset_line
670
671    for index, line in enumerate(script_lines, start = 1):
672        line = line.strip()
673        if line == "":
674            continue
675        if line.startswith("#"):
676            continue
677
678        sp = line.split(' ', 1)
679        cmd = sp[0]
680        if len(sp) == 2:
681            args = sp[1]
682        else:
683            args = ""
684
685        stateless_client.logger.log(format_text("Executing line {0} : '{1}'\n".format(index, line)))
686
687        if not cmd in cmd_table:
688            print("\n*** Error at line {0} : '{1}'\n".format(index, line))
689            stateless_client.logger.log(format_text("unknown command '{0}'\n".format(cmd), 'bold'))
690            return False
691
692        cmd_table[cmd](args)
693
694    stateless_client.logger.log(format_text("\n[Done]", 'bold'))
695
696    return True
697
698
699#
700def is_valid_file(filename):
701    if not os.path.isfile(filename):
702        raise argparse.ArgumentTypeError("The file '%s' does not exist" % filename)
703
704    return filename
705
706
707
708def setParserOptions():
709    parser = argparse.ArgumentParser(prog="trex_console.py")
710
711    parser.add_argument("-s", "--server", help = "TRex Server [default is localhost]",
712                        default = "localhost",
713                        type = str)
714
715    parser.add_argument("-p", "--port", help = "TRex Server Port  [default is 4501]\n",
716                        default = 4501,
717                        type = int)
718
719    parser.add_argument("--async_port", help = "TRex ASync Publisher Port [default is 4500]\n",
720                        default = 4500,
721                        dest='pub',
722                        type = int)
723
724    parser.add_argument("-u", "--user", help = "User Name  [default is currently logged in user]\n",
725                        default = get_current_user(),
726                        type = str)
727
728    parser.add_argument("-v", "--verbose", dest="verbose",
729                        action="store_true", help="Switch ON verbose option. Default is: OFF.",
730                        default = False)
731
732
733    parser.add_argument("--no_acquire", dest="acquire",
734                        action="store_false", help="Acquire all ports on connect. Default is: ON.",
735                        default = True)
736
737    parser.add_argument("--batch", dest="batch",
738                        nargs = 1,
739                        type = is_valid_file,
740                        help = "Run the console in a batch mode with file",
741                        default = None)
742
743    parser.add_argument("-t", "--tui", dest="tui",
744                        action="store_true", help="Starts with TUI mode",
745                        default = False)
746
747    parser.add_argument("-x", "--xtui", dest="xtui",
748                        action="store_true", help="Starts with XTERM TUI mode",
749                        default = False)
750
751    parser.add_argument("--top", dest="top",
752                        action="store_true", help="Set the window as always on top",
753                        default = False)
754
755    parser.add_argument("-q", "--quiet", dest="quiet",
756                        action="store_true", help="Starts with all outputs suppressed",
757                        default = False)
758
759    return parser
760
761
762def main():
763    parser = setParserOptions()
764    options = parser.parse_args()
765
766    if options.xtui:
767        options.tui = True
768
769    # always on top
770    if options.top:
771        set_window_always_on_top('trex_tui')
772
773
774    # Stateless client connection
775    if options.quiet:
776        verbose_level = LoggerApi.VERBOSE_QUIET
777    elif options.verbose:
778        verbose_level = LoggerApi.VERBOSE_HIGH
779    else:
780        verbose_level = LoggerApi.VERBOSE_REGULAR
781
782    # Stateless client connection
783    logger = ConsoleLogger()
784    stateless_client = STLClient(username = options.user,
785                                 server = options.server,
786                                 sync_port = options.port,
787                                 async_port = options.pub,
788                                 verbose_level = verbose_level,
789                                 logger = logger)
790
791    # TUI or no acquire will give us READ ONLY mode
792    try:
793        stateless_client.connect()
794    except STLError as e:
795        logger.log("Log:\n" + format_text(e.brief() + "\n", 'bold'))
796        return
797
798    if not options.tui and options.acquire:
799        try:
800            # acquire all ports
801            stateless_client.acquire()
802        except STLError as e:
803            logger.log("Log:\n" + format_text(e.brief() + "\n", 'bold'))
804            logger.log(format_text("\nSwitching to read only mode - only few commands will be available", 'bold'))
805
806
807    # a script mode
808    if options.batch:
809        cont = run_script_file(options.batch[0], stateless_client)
810        if not cont:
811            return
812
813    # console
814    try:
815        console = TRexConsole(stateless_client, options.verbose)
816        logger.prompt_redraw = console.prompt_redraw
817
818        # TUI
819        if options.tui:
820            console.do_tui("-x" if options.xtui else "")
821        else:
822            console.start()
823
824    except KeyboardInterrupt as e:
825        print("\n\n*** Caught Ctrl + C... Exiting...\n\n")
826
827    finally:
828        with stateless_client.logger.supress():
829            stateless_client.disconnect(stop_traffic = False)
830
831if __name__ == '__main__':
832
833    main()
834
835