src/python/srfsh.py

   1 #!/usr/bin/python
   2 # vim:et:ts=4
   3 """
   4 srfsh.py - provides a basic shell for issuing OpenSRF requests
   5
   6   help
   7     - show this menu
   8
   9   math_bench <count>
  10     - runs <count> opensrf.math requests and prints the average time
  11
  12   request <service> <method> [<param1>, <param2>, ...]
  13     - performs an opensrf request
  14
  15   set VAR=<value>
  16     - sets an environment variable
  17
  18   Environment variables:
  19     SRFSH_OUTPUT = pretty - print pretty JSON and key/value pairs for network objects
  20                  = raw - print formatted JSON
  21
  22     SRFSH_LOCALE = <locale> - request responses to be returned in locale <locale> if available
  23 """
  24
  25 import os, sys, time, readline, atexit, re
  26 import osrf.json, osrf.system, osrf.ses, osrf.conf, osrf.log, osrf.net
  27
  28 router_command_map = {
  29     'services' : 'opensrf.router.info.class.list',
  30     'service-stats' : 'opensrf.router.info.stats.class.node.all',
  31     'service-nodes' : 'opensrf.router.info.stats.class.all'
  32 }
  33
  34 # List of words to use for readline tab completion
  35 tab_complete_words = [
  36     'request',
  37     'set',
  38     'router',
  39     'help',
  40     'exit',
  41     'quit',
  42     'opensrf.settings',
  43     'opensrf.math'
  44 ]
  45
  46 # add the router commands to the tab-complete list
  47 for rcommand in router_command_map.keys():
  48     tab_complete_words.append(rcommand)
  49
  50 # -------------------------------------------------------------------
  51 # main listen loop
  52 # -------------------------------------------------------------------
  53 def do_loop():
  54
  55     command_map = {
  56         'request' : handle_request,
  57         'router' : handle_router,
  58         'math_bench' : handle_math_bench,
  59         'help' : handle_help,
  60         'set' : handle_set,
  61         'get' : handle_get,
  62     }
  63
  64     while True:
  65
  66         try:
  67             line = raw_input("srfsh# ")
  68
  69             if not len(line):
  70                 continue
  71
  72             if str.lower(line) == 'exit' or str.lower(line) == 'quit':
  73                 break
  74
  75             parts = str.split(line)
  76             command = parts.pop(0)
  77
  78             if command not in command_map:
  79                 print "unknown command: '%s'" % command
  80                 continue
  81
  82             command_map[command](parts)
  83
  84         except EOFError: # ^-d
  85             sys.exit(0)
  86
  87         except KeyboardInterrupt: # ^-c
  88             print ""
  89
  90         except Exception, e:
  91             print e
  92
  93
  94 def handle_router(parts):
  95
  96     if len(parts) == 0:
  97         print "usage: router <query>"
  98         return
  99
 100     query = parts[0]
 101
 102     if query not in router_command_map:
 103         print "router query options: %s" % ','.join(router_command_map.keys())
 104         return
 105
 106     return handle_request(['router', router_command_map[query]])
 107
 108 # -------------------------------------------------------------------
 109 # Set env variables to control behavior
 110 # -------------------------------------------------------------------
 111 def handle_set(parts):
 112     cmd = "".join(parts)
 113     pattern = re.compile('(.*)=(.*)').match(cmd)
 114     key = pattern.group(1)
 115     val = pattern.group(2)
 116     set_var(key, val)
 117     print "%s = %s" % (key, val)
 118
 119 def handle_get(parts):
 120     try:
 121         print get_var(parts[0])
 122     except:
 123         print ""
 124
 125
 126 # -------------------------------------------------------------------
 127 # Prints help info
 128 # -------------------------------------------------------------------
 129 def handle_help(foo):
 130     print __doc__
 131
 132 # -------------------------------------------------------------------
 133 # performs an opensrf request
 134 # -------------------------------------------------------------------
 135 def handle_request(parts):
 136
 137     if len(parts) < 2:
 138         print "usage: request <service> <api_name> [<param1>, <param2>, ...]"
 139         return
 140
 141     service = parts.pop(0)
 142     method = parts.pop(0)
 143     locale = __get_locale()
 144     jstr = '[%s]' % "".join(parts)
 145     params = None
 146
 147     try:
 148         params = osrf.json.to_object(jstr)
 149     except:
 150         print "Error parsing JSON: %s" % jstr
 151         return
 152
 153     ses = osrf.ses.ClientSession(service, locale=locale)
 154
 155     start = time.time()
 156
 157     req = ses.request2(method, tuple(params))
 158
 159
 160     while True:
 161         resp = None
 162         try:
 163             resp = req.recv(timeout=120)
 164         except osrf.net.XMPPNoRecipient:
 165             print "Unable to communicate with %s" % service
 166             total = 0
 167             break
 168
 169         osrf.log.log_internal("Looping through receive request")
 170         if not resp:
 171             break
 172         total = time.time() - start
 173
 174         otp = get_var('SRFSH_OUTPUT')
 175         if otp == 'pretty':
 176             print "Received Data: %s\n" % osrf.json.debug_net_object(resp.content())
 177         else:
 178             print "Received Data: %s\n" % osrf.json.pprint(osrf.json.to_json(resp.content()))
 179
 180     req.cleanup()
 181     ses.cleanup()
 182
 183     print '-'*60
 184     print "Total request time: %f" % total
 185     print '-'*60
 186
 187
 188 def handle_math_bench(parts):
 189
 190     count = int(parts.pop(0))
 191     ses = osrf.ses.ClientSession('opensrf.math')
 192     times = []
 193
 194     for cnt in range(100):
 195         if cnt % 10:
 196             sys.stdout.write('.')
 197         else:
 198             sys.stdout.write( str( cnt / 10 ) )
 199     print ""
 200
 201
 202     for cnt in range(count):
 203
 204         starttime = time.time()
 205         req = ses.request('add', 1, 2)
 206         resp = req.recv(timeout=2)
 207         endtime = time.time()
 208
 209         if resp.content() == 3:
 210             sys.stdout.write("+")
 211             sys.stdout.flush()
 212             times.append( endtime - starttime )
 213         else:
 214             print "What happened? %s" % str(resp.content())
 215
 216         req.cleanup()
 217         if not ( (cnt + 1) % 100):
 218             print ' [%d]' % (cnt + 1)
 219
 220     ses.cleanup()
 221     total = 0
 222     for cnt in times:
 223         total += cnt
 224     print "\naverage time %f" % (total / len(times))
 225
 226
 227
 228
 229 # -------------------------------------------------------------------
 230 # Defines the tab-completion handling and sets up the readline history
 231 # -------------------------------------------------------------------
 232 def setup_readline():
 233
 234     class SrfshCompleter(object):
 235
 236         def __init__(self, words):
 237             self.words = words
 238             self.prefix = None
 239
 240         def complete(self, prefix, index):
 241
 242             if prefix != self.prefix:
 243
 244                 self.prefix = prefix
 245
 246                 # find all words that start with this prefix
 247                 self.matching_words = [
 248                     w for w in self.words if w.startswith(prefix)
 249                 ]
 250
 251                 if len(self.matching_words) == 0:
 252                     return None
 253
 254                 if len(self.matching_words) == 1:
 255                     return self.matching_words[0]
 256
 257                 sys.stdout.write('\n%s\nsrfsh# %s' %
 258                     (' '.join(self.matching_words), readline.get_line_buffer()))
 259
 260                 return None
 261
 262     completer = SrfshCompleter(tuple(tab_complete_words))
 263     readline.parse_and_bind("tab: complete")
 264     readline.set_completer(completer.complete)
 265
 266     histfile = os.path.join(get_var('HOME'), ".srfsh_history")
 267     try:
 268         readline.read_history_file(histfile)
 269     except IOError:
 270         pass
 271     atexit.register(readline.write_history_file, histfile)
 272
 273     readline.set_completer_delims(readline.get_completer_delims().replace('-',''))
 274
 275
 276 def do_connect():
 277     file = os.path.join(get_var('HOME'), ".srfsh.xml")
 278     osrf.system.System.connect(config_file=file, config_context='srfsh')
 279
 280 def load_plugins():
 281     # Load the user defined external plugins
 282     # XXX Make this a real module interface, with tab-complete words, commands, etc.
 283     try:
 284         plugins = osrf.conf.get('plugins')
 285
 286     except:
 287         # XXX standard srfsh.xml does not yet define <plugins> element
 288         #print("No plugins defined in /srfsh/plugins/plugin\n")
 289         return
 290
 291     plugins = osrf.conf.get('plugins.plugin')
 292     if not isinstance(plugins, list):
 293         plugins = [plugins]
 294
 295     for module in plugins:
 296         name = module['module']
 297         init = module['init']
 298         print "Loading module %s..." % name
 299
 300         try:
 301             string = 'import %s\n%s.%s()' % (name, name, init)
 302             exec(string)
 303             print 'OK'
 304
 305         except Exception, e:
 306             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
 307
 308 def set_vars():
 309     if not get_var('SRFSH_OUTPUT'):
 310         set_var('SRFSH_OUTPUT', 'raw')
 311
 312     # XXX Do we need to differ between LANG and LC_MESSAGES?
 313     if not get_var('SRFSH_LOCALE'):
 314         set_var('SRFSH_LOCALE', get_var('LC_ALL'))
 315
 316 def set_var(key, val):
 317     os.environ[key] = val
 318
 319 def get_var(key):
 320     return os.environ.get(key, '')
 321
 322 def __get_locale():
 323     """
 324     Return the defined locale for this srfsh session.
 325
 326     A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern.
 327     This function munges the LC_ALL setting to conform to that pattern; for
 328     example, trimming en_CA.UTF-8 to en-CA.
 329
 330     >>> import srfsh
 331     >>> srfsh.set_var('SRFSH_LOCALE', 'zz-ZZ')
 332     >>> print __get_locale()
 333     zz-ZZ
 334     >>> srfsh.set_var('SRFSH_LOCALE', 'en_CA.UTF-8')
 335     >>> print __get_locale()
 336     en-CA
 337     """
 338
 339     env_locale = get_var('SRFSH_LOCALE')
 340     if env_locale:
 341         pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale)
 342         lang = pattern.group(1)
 343         region = pattern.group(2)
 344         locale = "%s-%s" % (lang, region)
 345     else:
 346         locale = 'en-US'
 347
 348     return locale
 349
 350 if __name__ == '__main__':
 351
 352     # Kick it off
 353     set_vars()
 354     setup_readline()
 355     do_connect()
 356     load_plugins()
 357     do_loop()
 358