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