src/target/trx_toolkit/ctrl_cmd.py

   1 #!/usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 # TRX Toolkit
   5 # Auxiliary tool to send custom commands via TRX CTRL interface,
   6 # which may be useful for testing and fuzzing
   7 #
   8 # (C) 2017-2018 by Vadim Yanitskiy <axilirator@gmail.com>
   9 #
  10 # All Rights Reserved
  11 #
  12 # This program is free software; you can redistribute it and/or modify
  13 # it under the terms of the GNU General Public License as published by
  14 # the Free Software Foundation; either version 2 of the License, or
  15 # (at your option) any later version.
  16 #
  17 # This program is distributed in the hope that it will be useful,
  18 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  19 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  20 # GNU General Public License for more details.
  21
  22 APP_CR_HOLDERS = [("2017-2018", "Vadim Yanitskiy <axilirator@gmail.com>")]
  23
  24 import logging as log
  25 import signal
  26 import argparse
  27 import select
  28 import sys
  29
  30 from app_common import ApplicationBase
  31 from udp_link import UDPLink
  32
  33 class Application(ApplicationBase):
  34         def __init__(self):
  35                 self.app_print_copyright(APP_CR_HOLDERS)
  36                 self.argv = self.parse_argv()
  37
  38                 # Set up signal handlers
  39                 signal.signal(signal.SIGINT, self.sig_handler)
  40
  41                 # Configure logging
  42                 self.app_init_logging(self.argv)
  43
  44                 # Init UDP connection
  45                 self.ctrl_link = UDPLink(
  46                         self.argv.remote_addr, self.argv.base_port + 1,
  47                         self.argv.bind_addr, self.argv.bind_port)
  48
  49                 # Debug print
  50                 log.info("Init TRXC interface (%s)" \
  51                         % self.ctrl_link.desc_link())
  52
  53         def parse_argv(self):
  54                 parser = argparse.ArgumentParser(prog = "ctrl_cmd",
  55                         description = "Auxiliary tool to send control commands")
  56
  57                 # Register common logging options
  58                 self.app_reg_logging_options(parser)
  59
  60                 trx_group = parser.add_argument_group("TRX interface")
  61                 trx_group.add_argument("-r", "--remote-addr",
  62                         dest = "remote_addr", type = str, default = "127.0.0.1",
  63                         help = "Set remote address (default %(default)s)")
  64                 trx_group.add_argument("-b", "--bind-addr",
  65                         dest = "bind_addr", type = str, default = "0.0.0.0",
  66                         help = "Set bind address (default %(default)s)")
  67                 trx_group.add_argument("-p", "--base-port",
  68                         dest = "base_port", type = int, default = 6700,
  69                         help = "Set base port number (default %(default)s)")
  70                 trx_group.add_argument("-P", "--bind-port",
  71                         dest = "bind_port", type = int, default = 0,
  72                         help = "Set bind port number (default random)")
  73                 trx_group.add_argument("-f", "--fuzzing",
  74                         dest = "fuzzing", action = "store_true",
  75                         help = "Send raw payloads (without CMD)")
  76
  77                 return parser.parse_args()
  78
  79         def run(self):
  80                 while True:
  81                         self.print_prompt()
  82
  83                         # Wait until we get any data on any socket
  84                         socks = [sys.stdin, self.ctrl_link.sock]
  85                         r_event, w_event, x_event = select.select(socks, [], [])
  86
  87                         # Check for incoming CTRL commands
  88                         if sys.stdin in r_event:
  89                                 cmd = sys.stdin.readline()
  90                                 self.handle_cmd(cmd)
  91
  92                         if self.ctrl_link.sock in r_event:
  93                                 data, addr = self.ctrl_link.sock.recvfrom(128)
  94                                 sys.stdout.write("\r%s\n" % data.decode())
  95                                 sys.stdout.flush()
  96
  97         def handle_cmd(self, cmd):
  98                 # Strip spaces, tabs, etc.
  99                 cmd = cmd.strip().strip("\0")
 100
 101                 # Send a command
 102                 if self.argv.fuzzing:
 103                         self.ctrl_link.send("%s" % cmd)
 104                 else:
 105                         self.ctrl_link.send("CMD %s\0" % cmd)
 106
 107         def print_prompt(self):
 108                 sys.stdout.write("TRXC# ")
 109                 sys.stdout.flush()
 110
 111         def sig_handler(self, signum, frame):
 112                 log.info("Signal %d received" % signum)
 113                 if signum == signal.SIGINT:
 114                         sys.exit(0)
 115
 116 if __name__ == '__main__':
 117         app = Application()
 118         app.run()