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                         from nodecommon import email_exception
 188                         email_exception()
 189                         return ("", "SCRIPTTIMEOUT")
 190
 191         def system(self, cmd, timeout=COMMAND_TIMEOUT*2):
 192                 (o,e) = self.run(cmd, timeout)
 193                 self.output = o
 194                 self.error = e
 195                 if self.s.returncode is None:
 196                         self.s.wait()
 197                 return self.s.returncode
 198
 199         def run(self, cmd, timeout=COMMAND_TIMEOUT*2):
 200
 201                 #print "CMD.run(%s)" % cmd
 202                 s = Sopen(cmd, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
 203                 self.s = s
 204                 (f_in, f_out, f_err) = (s.stdin, s.stdout, s.stderr)
 205                 #print "calling select(%s)" % timeout
 206                 lout, lin, lerr = select([f_out], [], [f_err], timeout)
 207                 #print "TIMEOUT!!!!!!!!!!!!!!!!!!!"
 208                 if len(lin) == 0 and len(lout) == 0 and len(lerr) == 0:
 209                         # Reached a timeout!  Nuke process so it does not hang.
 210                         #print "KILLING"
 211                         s.kill(signal.SIGKILL)
 212                         raise ExceptionTimeout("TIMEOUT Running: %s" % cmd)
 213                 else:
 214                         #print "RETURNING"
 215                         #print len(lin), len(lout), len(lerr)
 216                         pass
 217
 218                 o_value = ""
 219                 e_value = ""
 220
 221                 #print "reading from f_out"
 222                 if len(lout) > 0: o_value = f_out.read()
 223                 #print "reading from f_err"
 224                 if len(lerr) > 0: e_value = f_err.read()
 225
 226                 #print "striping output"
 227                 o_value = o_value.strip()
 228                 e_value = e_value.strip()
 229
 230                 #print "OUTPUT", o_value, e_value
 231
 232                 #print "closing files"
 233                 f_out.close()
 234                 f_in.close()
 235                 f_err.close()
 236                 try:
 237                         #print "s.kill()"
 238                         s.kill()
 239                         #print "after s.kill()"
 240                 except OSError:
 241                         # no such process, due to it already exiting...
 242                         pass
 243
 244                 #print o_value, e_value
 245                 return (o_value, e_value)
 246
 247         def runargs(self, args, timeout=COMMAND_TIMEOUT*2):
 248
 249                 #print "CMD.run(%s)" % " ".join(args)
 250                 s = Sopen(args, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
 251                 self.s = s
 252                 (f_in, f_out, f_err) = (s.stdin, s.stdout, s.stderr)
 253                 lout, lin, lerr = select([f_out], [], [f_err], timeout)
 254                 if len(lin) == 0 and len(lout) == 0 and len(lerr) == 0:
 255                         # Reached a timeout!  Nuke process so it does not hang.
 256                         s.kill(signal.SIGKILL)
 257                         raise ExceptionTimeout("TIMEOUT Running: %s" % cmd)
 258                 o_value = f_out.read()
 259                 e_value = ""
 260                 if o_value == "":       # An error has occured
 261                         e_value = f_err.read()
 262
 263                 o_value = o_value.strip()
 264                 e_value = e_value.strip()
 265
 266                 f_out.close()
 267                 f_in.close()
 268                 f_err.close()
 269                 try:
 270                         s.kill()
 271                 except OSError:
 272                         # no such process, due to it already exiting...
 273                         pass
 274
 275                 return (o_value, e_value)
 276
 277
 278 class SSH(CMD):
 279         def __init__(self, user, host, port=22, options = ssh_options):
 280                 self.options = options
 281                 self.user = user
 282                 self.host = host
 283                 self.port = port
 284                 return
 285
 286         def __options_to_str(self):
 287                 options = ""
 288                 for o,v in self.options.iteritems():
 289                         options = options + "-o %s=%s " % (o,v)
 290                 return options
 291
 292         def run(self, cmd, timeout=COMMAND_TIMEOUT*2):
 293                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 294                                                                         self.user, self.host, cmd)
 295                 #print "SSH.run(%s)" % cmd
 296                 return CMD.run(self, cmd, timeout)
 297
 298         def get_file(self, rmt_filename, local_filename=None):
 299                 if local_filename == None:
 300                         local_filename = "./"
 301                 cmd = "scp -P %s -B %s %s@%s:%s %s" % (self.port, self.__options_to_str(),
 302                                                                         self.user, self.host,
 303                                                                         rmt_filename, local_filename)
 304                 # output :
 305                 #       errors will be on stderr,
 306                 #   success will have a blank stderr...
 307                 return CMD.run_noexcept(self, cmd)
 308
 309         def run_noexcept(self, cmd):
 310                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 311                                                                         self.user, self.host, cmd)
 312                 #print "SSH.run_noexcept(%s)" % cmd
 313                 return CMD.run_noexcept(self, cmd)
 314
 315         def run_noexcept2(self, cmd, timeout=COMMAND_TIMEOUT*2):
 316                 cmd = "ssh -p %s %s %s@%s %s" % (self.port, self.__options_to_str(),
 317                                                                         self.user, self.host, cmd)
 318                 #print "SSH.run_noexcept2(%s)" % cmd
 319                 r = CMD.run_noexcept(self, cmd, timeout)
 320
 321                 # XXX: this may be resulting in deadlocks... not sure.
 322                 #if self.s.returncode is None:
 323                 #       #self.s.kill()
 324                 #       self.s.kill(signal.SIGKILL)
 325                 #       self.s.wait()
 326                 #       self.ret = self.s.returncode
 327                 self.ret = -1
 328
 329                 return r
 330
 331         def system2(self, cmd, timeout=COMMAND_TIMEOUT*2):
 332                 cmd = "ssh -p %s %s %s@%s %s" % (self.port, self.__options_to_str(),
 333                                                                         self.user, self.host, cmd)
 334                 #print "SSH.system2(%s)" % cmd
 335                 return CMD.system(self, cmd, timeout)
 336
 337         def runE(self, cmd):
 338                 cmd = "ssh -p %s %s %s@%s '%s'" % (self.port, self.__options_to_str(),
 339                                                                         self.user, self.host, cmd)
 340                 if ( DEBUG == 1 ):
 341                         print cmd,
 342                 (f_in, f_out, f_err) = os.popen3(cmd)
 343
 344                 value = f_out.read()
 345                 if value == "": # An error has occured
 346                         value = f_err.read()
 347                         value = value.strip()
 348
 349                 if ( DEBUG == 1 ):
 350                         print " == %s" % value
 351                 f_out.close()
 352                 f_in.close()
 353                 f_err.close()
 354                 return value.strip()
 355
 356 import time
 357 class MyTimer:
 358         def __init__(self):
 359                 self.start = time.time()
 360
 361         def end(self):
 362                 self.end = time.time()
 363                 t = self.end-self.start
 364                 return t
 365
 366         def diff(self):
 367                 self.end = time.time()
 368                 t = self.end-self.start
 369                 self.start = self.end
 370                 return t