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