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