DataExport.py

#!/usr/local/bin/python
#
#                               Copyright 2005
#                                     by
#                        The Board of Trustees of the
#                     Leland Stanford Junior University.
#                            All rights reserved.
#

__facility__ = "Online"
__abstract__ = "GLAST LAT data export routines"
__author__   = "J. Panetta <panetta@slac.stanford.edu> SLAC - GLAST LAT I&T/Online"
__date__     = ("$Date: 2006/01/31 17:24:49 $").split(' ')[1]
__version__  = "$Revision: 2.8 $"
__release__  = "$Name: R04-12-00 $"
__credits__  = "SLAC"

import LATTE.copyright_SLAC

import exceptions
import md5
import os
import shutil
import stat
import sys
import logging as log
from   LATTE.tools.FreeSpace        import FreeSpace
from   LATTE.client.gOptions        import Options

class GExportException(exceptions.Exception):
  """!\brief Define an exception from the Export system.
  """
  def __init__(self, status):
    self.status = status

  def __str__(self):
    return repr(self.status)


class DataExport:
  """!\brief Class to synch a file or set of files between one location and another.
  """
  # Error codes
  NO_ERROR                      = 0x0
  DIR_NOT_WRITABLE              = 0x1
  NOT_EXTANT                    = 0x2

  DISK_FULL                     = 0x10

  DISK_FULL_LEVEL               = 100*1024*1024  # 100 Megabytes

  def __init__(self):
    self.__purgeFiles   = False

    self.transfer = self.__copy   # Other valid is self.__rename
    pass

  def setPurge(self, purge):
    """
    """
    if purge == True:
      self.__purgeFiles = True
    else:
      self.__purgeFiles = False

  def synch(self, source, dest):
    """!\brief Synch up two directory structures.

       Synch only goes in one direction:  source --> destination.
    """

    sourceStatus = self.__checkDir(source)
    destStatus   = self.__checkDir(dest)

    # do something with statuses

    # get a list of directories under source.  Do this only once, as it can take time
    dirList = []
    for file in os.listdir(source):
      file = os.path.join(source,file)
      if os.path.isdir(file):
        dirList.append(file)

    for sourceDir in dirList:
      destDir = os.path.join(dest, os.path.basename(sourceDir))
      try:
        log.info( "Exporting from %s to %s" % ( sourceDir, destDir ) )
        self.exportDir(sourceDir, destDir)
      except Exception, e:
        log.error( "Failed exporting from %s to %s" % ( sourceDir, destDir ) )
        log.error( "  Exception: %s" % e )
        pass

  def export(self, source, dest):
    """!\brief Copy from source to dest.

       This routine will determine whether source is a file or a directory
    """

    self.__checkDirs(source, dest)

    # Check for common disk on Windows XP/NT (*not* CE).
    if os.name == 'nt':
      commonPath = os.path.commonprefix( [source, dest] )
      if len(commonPath) > 0 and self.__purgeFiles:
        self.transfer = self.__rename
      else:
        self.transfer = self.__copy
    else:
      pass # self.__transfer is set in constructor to self.__copy

    self.transfer(source, dest)

  def md5sum(self, fname):
    """!\brief Returns an md5 hash for file fname.
    """
    # md5hash from ActiveState cookbook
    # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/266486
    #
    # 2004-01-30
    #
    # Nick Vargish
    #
    # Simple md5 hash utility for generating md5 checksums of files.
    # Modified slightly by Jim Panetta
    try:
      if os.path.isdir(fname):
        return "Directory"
      f = file(fname, 'rb')
    except:
      return 'Failed to open file'
    ret = self.__sumfile(f)
    f.close()
    return ret


  def __copy(self, source, dest):
    # copy/checksum/delete
    if os.path.isdir(source):
      shutil.copytree( source, dest )      # copy
      self.__compareTree( source, dest )   # checksum

      if self.__purgeFiles:                # delete
        for root, dirs, files in os.walk(source, topdown=False):
          for name in files:
              f = os.path.join(root, name)
              os.chmod(f, stat.S_IWRITE) # Just in case it is read-only
              os.remove(f)
          for name in dirs:
              os.rmdir(os.path.join(root, name))
        os.rmdir(root)
    else:  # regular file.  much simpler
      shutil.copy( source, dest )         # copy
      sourceMD5 = self.md5sum(source)     # checksum
      destMD5   = self.md5sum(dest)
      if sourceMD5 != destMD5:
        msg = "md5 sums don't match for %s :" % (source)
        msg += "\n   source: %s, dest: %s" % ( sourceMD5, destMD5 )
        raise GExportException, msg

      if self.__purgeFiles:               # delete
        os.remove(source)

  def __rename(self, source, dest):
    # just move it.  No purging needed
    shutil.move(source, dest)

  def __checkDirs(self, source, dest):
    # check statuses
    sourceStatus = self.__checkStat(source)
    destStatus   = self.__checkStat(dest)

    # don't care if source dir is on a full disk or if it isn't writable
    if ( sourceStatus & ~(DataExport.DISK_FULL|DataExport.DIR_NOT_WRITABLE) ):
      msg = "Export Error: Attempt to read from source %s encountered errors: " % (source)
      msg += self.__statusParse(sourceStatus)
      raise GExportException, msg
    # Destination must not be extant.
    if not ( destStatus & DataExport.NOT_EXTANT ):
      msg = "Export Error: Destination %s already exists" % (dest)
      raise GExportException, msg
    if ( destStatus & ~DataExport.NOT_EXTANT ) :
      msg = "Export Error: Attempt to write to destination %s encountered errors: " % (dest)
      msg += self.__statusParse(destStatus)

      raise GExportException, msg
    # source and dest can't be the same
    if os.path.abspath(source) == os.path.abspath(dest):
      msg = "Source and destination directories are the same.  Export aborted"
      raise GExportException, msg


  def __compareTree(self, sourceDir, destDir):
    """!\brief Compare all files under sourceDir and destDir with md5.

    \param sourceDir Source directory
    \param destDir   Destination directory
     """
    # scan the source/dest directories
    # chdir so that the source/dest keys are the same
    saveDir = os.getcwd()
    os.chdir(sourceDir)
    sourceWalk = os.walk(".")
    sourceMD5 = self.__treeWalk( sourceWalk, self.md5sum )

    os.chdir(destDir)
    destWalk   = os.walk(".")
    destMD5   = self.__treeWalk( destWalk,   self.md5sum )

    # compare the checksums
    for file in sourceMD5.keys():
      if destMD5.has_key(file):
        if sourceMD5[file] != destMD5[file]:
          os.chdir(saveDir)
          msg = "md5 sums don't match for %s :" % (file)
          msg += "\n   source: %s, dest: %s" % ( sourceMD5[file], destMD5[file])
          raise GExportException, msg
        else:
          destMD5.pop(file)
      else:
        os.chdir(saveDir)
        msg = "Destination directory %s does not contain file %s from source directory %s" % \
              ( destDir, file, sourceDir )
        raise GExportException, msg

    if destMD5 != {}:
      os.chdir(saveDir)
      msg = "Destination directory %s contains files not in source directory:" % ( destDir )
      for file in destMD5.keys():
        msg += "\n    %s" % ( file )
      raise GExportException, msg
    os.chdir(saveDir)

  def __treeWalk(self, tree, function):
    """!\brief Walk a tree, executing function() on every file.
    """
    returnDict = {}
    for root, dirs, files in tree:
      for name in dirs + files:
        try:
          retVal = function(os.path.join(root, name))
        except:
          retVal = -1
        returnDict[os.path.join(root, name)] = retVal
    return returnDict

  def __checkStat(self, obj):
    """!\brief Check if source and directories exist, are writable,
    and the disks aren't full.

    \param obj File path
    \return Status
    """
    status = 0L

    if not os.path.exists(obj):
      status |= DataExport.NOT_EXTANT
    elif FreeSpace(os.path.dirname(obj)) < DataExport.DISK_FULL_LEVEL:
      status |= DataExport.DISK_FULL
    if not os.access(os.path.dirname(obj),os.W_OK):
      status |= DataExport.DIR_NOT_WRITABLE

    return status

  def __makeReadOnly(self, file):
    status = 0
    try:
      os.chmod(file, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
    except:
      status = -1
    return status

  def __statusParse(self, status):
    msg = ""
    if status & DataExport.DIR_NOT_WRITABLE:
      msg += "\n  Directory not writable"
    if status & DataExport.NOT_EXTANT:
      msg += "\n  Directory does not exist"
    if status & DataExport.DISK_FULL:
      msg += "\n  Disk is full"
    return msg

  def __sumfile(self, fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = md5.new()
    while True:
      d = fobj.read(8096)
      if not d:
        break
      m.update(d)
    return m.hexdigest()



def usage():
  return """
DataExport:  Securely copy files from one directory to another.

Mandatory options:
--source        Source directory
--destination   Destination directory

Flags:
--purge         Set if source files are to be deleted after copy.
"""

if __name__ == "__main__":
  """Command line interface for the Data Export system.
     Can be called by the user or by an automated process.

  """

  # Options, first list mandatory, second list optional
  options = Options(['source', 'destination'], [], ['purge'])
  try:
    options.parse()
  except Exception, msg:
    options.usage(usage())
    raise Exception, msg

  source         = options.source
  destination    = options.destination
  print "source=%s, dest=%s" % ( source, destination)

  exporter = DataExport()

  # Do we want file deletion?
  if options.purge is not None:
    exporter.setPurge(True)

  exporter.export(source, destination)

---