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