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