soltesz.py

   1 import os
   2 import sys
   3 import pickle
   4 noserial=False
   5 try:
   6         from PHPSerialize import *
   7         from PHPUnserialize import *
   8 except:
   9         #print >>sys.stderr, "PHPSerial db type not allowed."
  10         noserial=True
  11
  12 import inspect
  13 import shutil
  14 import config
  15
  16 import config
  17
  18 DEBUG= 0
  19 PICKLE_PATH=config.MONITOR_DATA_ROOT
  20
  21 class ExceptionTimeout(Exception): pass
  22
  23 def dbLoad(name, type=None):
  24         return SPickle().load(name, type)
  25
  26 def dbExists(name, type=None):
  27         #if self.config.debug:
  28         #       name = "debug.%s" % name
  29         return SPickle().exists(name, type)
  30
  31 def dbDump(name, obj=None, type=None):
  32         # depth of the dump is 2 now, since we're redirecting to '.dump'
  33         return SPickle().dump(name, obj, type, 2)
  34
  35 def if_cached_else_refresh(cond, refresh, name, function, type=None):
  36         s = SPickle()
  37         if refresh:
  38                 if not config.debug and s.exists("production.%s" % name, type):
  39                         s.remove("production.%s" % name, type)
  40                 if config.debug and s.exists("debug.%s" % name, type):
  41                         s.remove("debug.%s" % name, type)
  42
  43         return if_cached_else(cond, name, function, type)
  44
  45 def if_cached_else(cond, name, function, type=None):
  46         s = SPickle()
  47         if (cond and s.exists("production.%s" % name, type)) or \
  48            (cond and config.debug and s.exists("debug.%s" % name, type)):
  49                 o = s.load(name, type)
  50         else:
  51                 o = function()
  52                 if cond:
  53                         s.dump(name, o, type)   # cache the object using 'name'
  54                         o = s.load(name, type)
  55                 # TODO: what if 'o' hasn't been converted...
  56         return o
  57
  58 class SPickle:
  59         def __init__(self, path=PICKLE_PATH):
  60                 self.path = path
  61
  62         def if_cached_else(self, cond, name, function, type=None):
  63                 if cond and self.exists("production.%s" % name, type):
  64                         o = self.load(name, type)
  65                 else:
  66                         o = function()
  67                         if cond:
  68                                 self.dump(name, o, type)        # cache the object using 'name'
  69                 return o
  70
  71         def __file(self, name, type=None):
  72                 if type == None:
  73                         return "%s/%s.pkl" % (self.path, name)
  74                 else:
  75                         if noserial:
  76                                 raise Exception("No PHPSerializer module available")
  77
  78                         return "%s/%s.phpserial" % (self.path, name)
  79
  80         def exists(self, name, type=None):
  81                 return os.path.exists(self.__file(name, type))
  82
  83         def remove(self, name, type=None):
  84                 return os.remove(self.__file(name, type))
  85
  86         def load(self, name, type=None):
  87                 """
  88                 In debug mode, we should fail if neither file exists.
  89                         if the debug file exists, reset name
  90                         elif the original file exists, make a copy, reset name
  91                         else neither exist, raise an error
  92                 Otherwise, it's normal mode, if the file doesn't exist, raise error
  93                 Load the file
  94                 """
  95
  96                 if config.debug:
  97                         if self.exists("debug.%s" % name, type):
  98                                 name = "debug.%s" % name
  99                         elif self.exists("production.%s" % name, type):
 100                                 debugname = "debug.%s" % name
 101                                 if not self.exists(debugname, type):
 102                                         name = "production.%s" % name
 103                                         shutil.copyfile(self.__file(name, type), self.__file(debugname, type))
 104                                 name = debugname
 105                         else:   # neither exist
 106                                 raise Exception, "No such pickle based on %s" % self.__file("debug.%s" % name, type)
 107                 else:
 108                         if   self.exists("production.%s" % name, type):
 109                                 name = "production.%s" % name
 110                         elif self.exists(name, type):
 111                                 name = name
 112                         else:
 113                                 raise Exception, "No such file %s" % name
 114
 115
 116                 #print "loading %s" % self.__file(name, type)
 117                 f = open(self.__file(name, type), 'r')
 118                 if type == None:
 119                         o = pickle.load(f)
 120                 else:
 121                         if noserial:
 122                                 raise Exception("No PHPSerializer module available")
 123                         s = PHPUnserialize()
 124                         o = s.unserialize(f.read())
 125                 f.close()
 126                 return o
 127
 128
 129         # use the environment to extract the data associated with the local
 130         # variable 'name'
 131         def dump(self, name, obj=None, type=None, depth=1):
 132                 if obj == None:
 133                         o = inspect.getouterframes(inspect.currentframe())
 134                         up1 = o[depth][0] # get the frame one prior to (up from) this frame
 135                         argvals = inspect.getargvalues(up1)
 136                         # TODO: check that 'name' is a local variable; otherwise this would fail.
 137                         obj = argvals[3][name] # extract the local variable name 'name'
 138                 if not os.path.isdir("%s/" % self.path):
 139                         os.mkdir("%s" % self.path)
 140                 if config.debug:
 141                         name = "debug.%s" % name
 142                 else:
 143                         name = "production.%s" % name
 144                 f = open(self.__file(name, type), 'w')
 145                 if type == None:
 146                         pickle.dump(obj, f)
 147                 else:
 148                         if noserial:
 149                                 raise Exception("No PHPSerializer module available")
 150                         s = PHPSerialize()
 151                         f.write(s.serialize(obj))
 152                 f.close()
 153                 return
 154
 155
 156 COMMAND_TIMEOUT = 60
 157 ssh_options = { 'StrictHostKeyChecking':'no',
 158                                 'BatchMode':'yes',
 159                                 'PasswordAuthentication':'no',
 160                                 'ConnectTimeout':'%s' % COMMAND_TIMEOUT}
 161 from select import select
 162 import subprocess
 163 import signal
 164
 165 class Sopen(subprocess.Popen):
 166         def kill(self, signal = signal.SIGTERM):
 167                 os.kill(self.pid, signal)
 168
 169 def read_t(stream, count, timeout=COMMAND_TIMEOUT*2):
 170         lin, lout, lerr = select([stream], [], [], timeout)
 171         if len(lin) == 0:
 172                 raise ExceptionTimeout("TIMEOUT Running: %s" % cmd)
 173
 174         return stream.read(count)
 175
 176 class CMD:
 177         def __init__(self):
 178                 pass
 179
 180         def run_noexcept(self, cmd, timeout=COMMAND_TIMEOUT*2):
 181
 182                 #print "CMD.run_noexcept(%s)" % cmd
 183                 try:
 184                         return CMD.run(self,cmd,timeout)
 185                 except ExceptionTimeout:
 186                         import traceback; print traceback.print_exc()
 187                         return ("", "SCRIPTTIMEOUT")
 188
 189         def system(self, cmd, timeout=COMMAND_TIMEOUT*2):
 190                 (o,e) = self.run(cmd, timeout)
 191                 self.output = o
 192                 self.error = e
 193                 if self.s.returncode is None:
 194                         self.s.wait()
 195                 return self.s.returncode
 196
 197         def run(self, cmd, timeout=COMMAND_TIMEOUT*2):
 198
 199                 #print "CMD.run(%s)" % cmd
 200                 s = Sopen(cmd, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
 201                 self.s = s
 202                 (f_in, f_out, f_err) = (s.stdin, s.stdout, s.stderr)
 203                 #print "calling select(%s)" % timeout
 204                 lout, lin, lerr = select([f_out], [], [f_err], timeout)
 205                 #print "TIMEOUT!!!!!!!!!!!!!!!!!!!"
 206                 if len(lin) == 0 and len(lout) == 0 and len(lerr) == 0:
 207                         # Reached a timeout!  Nuke process so it does not hang.
 208                         #print "KILLING"
 209                         s.kill(signal.SIGKILL)
 210                         raise ExceptionTimeout("TIMEOUT Running: %s" % cmd)
 211                 else:
 212                         #print "RETURNING"
 213                         #print len(lin), len(lout), len(lerr)
 214                         pass
 215
 216                 o_value = ""
 217                 e_value = ""
 218
 219                 #print "reading from f_out"
 220                 if len(lout) > 0: o_value = f_out.read()
 221                 #print "reading from f_err"
 222                 if len(lerr) > 0: e_value = f_err.read()
 223
 224                 #print "striping output"
 225                 o_value = o_value.strip()
 226                 e_value = e_value.strip()
 227
 228                 #print "OUTPUT", o_value, e_value
 229
 230                 #print "closing files"
 231                 f_out.close()
 232                 f_in.close()
 233                 f_err.close()
 234                 try:
 235                         #print "s.kill()"
 236                         s.kill()
 237                         #print "after s.kill()"
 238                 except OSError:
 239                         # no such process, due to it already exiting...
 240                         pass
 241
 242                 #print o_value, e_value
 243                 return (o_value, e_value)
 244
 245         def runargs(self, args, timeout=COMMAND_TIMEOUT*2):
 246
 247                 #print "CMD.run(%s)" % " ".join(args)
 248                 s = Sopen(args, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
 249                 self.s = s
 250                 (f_in, f_out, f_err) = (s.stdin, s.stdout, s.stderr)
 251                 lout, lin, lerr = select([f_out], [], [f_err], timeout)
 252                 if len(lin) == 0 and len(lout) == 0 and len(lerr) == 0:
 253                         # Reached a timeout!  Nuke process so it does not hang.
 254                         s.kill(signal.SIGKILL)
 255                         raise ExceptionTimeout("TIMEOUT Running: %s" % cmd)
 256                 o_value = f_out.read()
 257                 e_value = ""
 258                 if o_value == "":       # An error has occured
 259                         e_value = f_err.read()
 260
 261                 o_value = o_value.strip()
 262                 e_value = e_value.strip()
 263
 264                 f_out.close()
 265                 f_in.close()
 266                 f_err.close()
 267                 try:
 268                         s.kill()
 269                 except OSError:
 270                         # no such process, due to it already exiting...
 271                         pass
 272
 273                 return (o_value, e_value)
 274
 275
 276 class SSH(CMD):
 277         def __init__(self, user, host, port=22, options = ssh_options):
 278                 self.options = options
 279                 self.user = user
 280                 self.host = host
 281                 self.port = port
 282                 return
 283
 284         def __options_to_str(self):
 285                 options = ""
 286                 for o,v in self.options.iteritems():
 287                         options = options + "-o %s=%s " % (o,v)
 288                 return options
 289
 290         def run(self, cmd, timeout=COMMAND_TIMEOUT*2):
 291                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 292                                                                         self.user, self.host, cmd)
 293                 #print "SSH.run(%s)" % cmd
 294                 return CMD.run(self, cmd, timeout)
 295
 296         def get_file(self, rmt_filename, local_filename=None):
 297                 if local_filename == None:
 298                         local_filename = "./"
 299                 cmd = "scp -P %s -B %s %s@%s:%s %s" % (self.port, self.__options_to_str(),
 300                                                                         self.user, self.host,
 301                                                                         rmt_filename, local_filename)
 302                 # output :
 303                 #       errors will be on stderr,
 304                 #   success will have a blank stderr...
 305                 return CMD.run_noexcept(self, cmd)
 306
 307         def run_noexcept(self, cmd):
 308                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 309                                                                         self.user, self.host, cmd)
 310                 #print "SSH.run_noexcept(%s)" % cmd
 311                 return CMD.run_noexcept(self, cmd)
 312
 313         def run_noexcept2(self, cmd, timeout=COMMAND_TIMEOUT*2):
 314                 cmd = "ssh -p %s %s %s@%s %s" % (self.port, self.__options_to_str(),
 315                                                                         self.user, self.host, cmd)
 316                 #print "SSH.run_noexcept2(%s)" % cmd
 317                 r = CMD.run_noexcept(self, cmd, timeout)
 318
 319                 # XXX: this may be resulting in deadlocks... not sure.
 320                 #if self.s.returncode is None:
 321                 #       #self.s.kill()
 322                 #       self.s.kill(signal.SIGKILL)
 323                 #       self.s.wait()
 324                 #       self.ret = self.s.returncode
 325                 self.ret = -1
 326
 327                 return r
 328
 329         def system2(self, cmd, timeout=COMMAND_TIMEOUT*2):
 330                 cmd = "ssh -p %s %s %s@%s %s" % (self.port, self.__options_to_str(),
 331                                                                         self.user, self.host, cmd)
 332                 #print "SSH.system2(%s)" % cmd
 333                 return CMD.system(self, cmd, timeout)
 334
 335         def runE(self, cmd):
 336                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 337                                                                         self.user, self.host, cmd)
 338                 if ( DEBUG == 1 ):
 339                         print cmd,
 340                 (f_in, f_out, f_err) = os.popen3(cmd)
 341
 342                 value = f_out.read()
 343                 if value == "": # An error has occured
 344                         value = f_err.read()
 345                         value = value.strip()
 346
 347                 if ( DEBUG == 1 ):
 348                         print " == %s" % value
 349                 f_out.close()
 350                 f_in.close()
 351                 f_err.close()
 352                 return value.strip()
 353
 354 import time
 355 class MyTimer:
 356         def __init__(self):
 357                 self.start = time.time()
 358
 359         def end(self):
 360                 self.end = time.time()
 361                 t = self.end-self.start
 362                 return t
 363
 364         def diff(self):
 365                 self.end = time.time()
 366                 t = self.end-self.start
 367                 self.start = self.end
 368                 return t