src/allmydata/util/fileutil.py

   1 """
   2 Futz with files like a pro.
   3 """
   4
   5 import sys, exceptions, os, stat, tempfile, time, binascii
   6
   7 from twisted.python import log
   8
   9 from pycryptopp.cipher.aes import AES
  10
  11
  12 def rename(src, dst, tries=4, basedelay=0.1):
  13     """ Here is a superkludge to workaround the fact that occasionally on
  14     Windows some other process (e.g. an anti-virus scanner, a local search
  15     engine, etc.) is looking at your file when you want to delete or move it,
  16     and hence you can't.  The horrible workaround is to sit and spin, trying
  17     to delete it, for a short time and then give up.
  18
  19     With the default values of tries and basedelay this can block for less
  20     than a second.
  21
  22     @param tries: number of tries -- each time after the first we wait twice
  23     as long as the previous wait
  24     @param basedelay: how long to wait before the second try
  25     """
  26     for i in range(tries-1):
  27         try:
  28             return os.rename(src, dst)
  29         except EnvironmentError, le:
  30             # XXX Tighten this to check if this is a permission denied error (possibly due to another Windows process having the file open and execute the superkludge only in this case.
  31             log.msg("XXX KLUDGE Attempting to move file %s => %s; got %s; sleeping %s seconds" % (src, dst, le, basedelay,))
  32             time.sleep(basedelay)
  33             basedelay *= 2
  34     return os.rename(src, dst) # The last try.
  35
  36 def remove(f, tries=4, basedelay=0.1):
  37     """ Here is a superkludge to workaround the fact that occasionally on
  38     Windows some other process (e.g. an anti-virus scanner, a local search
  39     engine, etc.) is looking at your file when you want to delete or move it,
  40     and hence you can't.  The horrible workaround is to sit and spin, trying
  41     to delete it, for a short time and then give up.
  42
  43     With the default values of tries and basedelay this can block for less
  44     than a second.
  45
  46     @param tries: number of tries -- each time after the first we wait twice
  47     as long as the previous wait
  48     @param basedelay: how long to wait before the second try
  49     """
  50     try:
  51         os.chmod(f, stat.S_IWRITE | stat.S_IEXEC | stat.S_IREAD)
  52     except:
  53         pass
  54     for i in range(tries-1):
  55         try:
  56             return os.remove(f)
  57         except EnvironmentError, le:
  58             # XXX Tighten this to check if this is a permission denied error (possibly due to another Windows process having the file open and execute the superkludge only in this case.
  59             if not os.path.exists(f):
  60                 return
  61             log.msg("XXX KLUDGE Attempting to remove file %s; got %s; sleeping %s seconds" % (f, le, basedelay,))
  62             time.sleep(basedelay)
  63             basedelay *= 2
  64     return os.remove(f) # The last try.
  65
  66 class ReopenableNamedTemporaryFile:
  67     """
  68     This uses tempfile.mkstemp() to generate a secure temp file.  It then closes
  69     the file, leaving a zero-length file as a placeholder.  You can get the
  70     filename with ReopenableNamedTemporaryFile.name.  When the
  71     ReopenableNamedTemporaryFile instance is garbage collected or its shutdown()
  72     method is called, it deletes the file.
  73     """
  74     def __init__(self, *args, **kwargs):
  75         fd, self.name = tempfile.mkstemp(*args, **kwargs)
  76         os.close(fd)
  77
  78     def __repr__(self):
  79         return "<%s instance at %x %s>" % (self.__class__.__name__, id(self), self.name)
  80
  81     def __str__(self):
  82         return self.__repr__()
  83
  84     def __del__(self):
  85         self.shutdown()
  86
  87     def shutdown(self):
  88         remove(self.name)
  89
  90 class NamedTemporaryDirectory:
  91     """
  92     This calls tempfile.mkdtemp(), stores the name of the dir in
  93     self.name, and rmrf's the dir when it gets garbage collected or
  94     "shutdown()".
  95     """
  96     def __init__(self, cleanup=True, *args, **kwargs):
  97         """ If cleanup, then the directory will be rmrf'ed when the object is shutdown. """
  98         self.cleanup = cleanup
  99         self.name = tempfile.mkdtemp(*args, **kwargs)
 100
 101     def __repr__(self):
 102         return "<%s instance at %x %s>" % (self.__class__.__name__, id(self), self.name)
 103
 104     def __str__(self):
 105         return self.__repr__()
 106
 107     def __del__(self):
 108         try:
 109             self.shutdown()
 110         except:
 111             import traceback
 112             traceback.print_exc()
 113
 114     def shutdown(self):
 115         if self.cleanup and hasattr(self, 'name'):
 116             rm_dir(self.name)
 117
 118 class EncryptedTemporaryFile:
 119     # not implemented: next, readline, readlines, xreadlines, writelines
 120
 121     def __init__(self):
 122         self.file = tempfile.TemporaryFile()
 123         self.key = os.urandom(16)  # AES-128
 124
 125     def _crypt(self, offset, data):
 126         offset_big = offset // 16
 127         offset_small = offset % 16
 128         iv = binascii.unhexlify("%032x" % offset_big)
 129         cipher = AES(self.key, iv=iv)
 130         cipher.process("\x00"*offset_small)
 131         return cipher.process(data)
 132
 133     def close(self):
 134         self.file.close()
 135
 136     def flush(self):
 137         self.file.flush()
 138
 139     def seek(self, offset, whence=0):  # 0 = SEEK_SET
 140         self.file.seek(offset, whence)
 141
 142     def tell(self):
 143         offset = self.file.tell()
 144         return offset
 145
 146     def read(self, size=-1):
 147         index = self.file.tell()
 148         ciphertext = self.file.read(size)
 149         plaintext = self._crypt(index, ciphertext)
 150         return plaintext
 151
 152     def write(self, plaintext):
 153         index = self.file.tell()
 154         ciphertext = self._crypt(index, plaintext)
 155         self.file.write(ciphertext)
 156
 157     def truncate(self, newsize):
 158         self.file.truncate(newsize)
 159
 160
 161 def make_dirs(dirname, mode=0777):
 162     """
 163     An idempotent version of os.makedirs().  If the dir already exists, do
 164     nothing and return without raising an exception.  If this call creates the
 165     dir, return without raising an exception.  If there is an error that
 166     prevents creation or if the directory gets deleted after make_dirs() creates
 167     it and before make_dirs() checks that it exists, raise an exception.
 168     """
 169     tx = None
 170     try:
 171         os.makedirs(dirname, mode)
 172     except OSError, x:
 173         tx = x
 174
 175     if not os.path.isdir(dirname):
 176         if tx:
 177             raise tx
 178         raise exceptions.IOError, "unknown error prevented creation of directory, or deleted the directory immediately after creation: %s" % dirname # careful not to construct an IOError with a 2-tuple, as that has a special meaning...
 179
 180 def rm_dir(dirname):
 181     """
 182     A threadsafe and idempotent version of shutil.rmtree().  If the dir is
 183     already gone, do nothing and return without raising an exception.  If this
 184     call removes the dir, return without raising an exception.  If there is an
 185     error that prevents deletion or if the directory gets created again after
 186     rm_dir() deletes it and before rm_dir() checks that it is gone, raise an
 187     exception.
 188     """
 189     excs = []
 190     try:
 191         os.chmod(dirname, stat.S_IWRITE | stat.S_IEXEC | stat.S_IREAD)
 192         for f in os.listdir(dirname):
 193             fullname = os.path.join(dirname, f)
 194             if os.path.isdir(fullname):
 195                 rm_dir(fullname)
 196             else:
 197                 remove(fullname)
 198         os.rmdir(dirname)
 199     except Exception, le:
 200         # Ignore "No such file or directory"
 201         if (not isinstance(le, OSError)) or le.args[0] != 2:
 202             excs.append(le)
 203
 204     # Okay, now we've recursively removed everything, ignoring any "No
 205     # such file or directory" errors, and collecting any other errors.
 206
 207     if os.path.exists(dirname):
 208         if len(excs) == 1:
 209             raise excs[0]
 210         if len(excs) == 0:
 211             raise OSError, "Failed to remove dir for unknown reason."
 212         raise OSError, excs
 213
 214
 215 def remove_if_possible(f):
 216     try:
 217         remove(f)
 218     except:
 219         pass
 220
 221 def open_or_create(fname, binarymode=True):
 222     try:
 223         return open(fname, binarymode and "r+b" or "r+")
 224     except EnvironmentError:
 225         return open(fname, binarymode and "w+b" or "w+")
 226
 227 def du(basedir):
 228     size = 0
 229
 230     for root, dirs, files in os.walk(basedir):
 231         for f in files:
 232             fn = os.path.join(root, f)
 233             size += os.path.getsize(fn)
 234
 235     return size
 236
 237 def move_into_place(source, dest):
 238     """Atomically replace a file, or as near to it as the platform allows.
 239     The dest file may or may not exist."""
 240     if "win32" in sys.platform.lower():
 241         remove_if_possible(dest)
 242     os.rename(source, dest)
 243
 244 def write(path, data):
 245     wf = open(path, "wb")
 246     try:
 247         wf.write(data)
 248     finally:
 249         wf.close()
 250
 251 def read(path):
 252     rf = open(path, "rb")
 253     try:
 254         return rf.read()
 255     finally:
 256         rf.close()
 257
 258 def put_file(pathname, inf):
 259     # TODO: create temporary file and move into place?
 260     outf = open(os.path.expanduser(pathname), "wb")
 261     try:
 262         while True:
 263             data = inf.read(32768)
 264             if not data:
 265                 break
 266             outf.write(data)
 267     finally:
 268         outf.close()