BaseTools/Source/Python/UPT/Core/FileHook.py

   1 ## @file
   2 # This file hooks file and directory creation and removal
   3 #
   4 # Copyright (c) 2014, Intel Corporation. All rights reserved.<BR>
   5 #
   6 # This program and the accompanying materials are licensed and made available
   7 # under the terms and conditions of the BSD License which accompanies this
   8 # distribution. The full text of the license may be found at
   9 # http://opensource.org/licenses/bsd-license.php
  10 #
  11 # THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  13 #
  14
  15 '''
  16 File hook
  17 '''
  18
  19 import os
  20 import stat
  21 import time
  22 import zipfile
  23 from time import sleep
  24 from Library import GlobalData
  25
  26 __built_in_remove__ = os.remove
  27 __built_in_mkdir__  = os.mkdir
  28 __built_in_rmdir__  = os.rmdir
  29 __built_in_chmod__  = os.chmod
  30 __built_in_open__   = open
  31
  32 _RMFILE      = 0
  33 _MKFILE      = 1
  34 _RMDIR       = 2
  35 _MKDIR       = 3
  36 _CHMOD       = 4
  37
  38 gBACKUPFILE = 'file.backup'
  39 gEXCEPTION_LIST = ['Conf'+os.sep+'DistributionPackageDatabase.db', '.tmp', gBACKUPFILE]
  40
  41 class _PathInfo:
  42     def __init__(self, action, path, mode=-1):
  43         self.action = action
  44         self.path = path
  45         self.mode = mode
  46
  47 class RecoverMgr:
  48     def __init__(self, workspace):
  49         self.rlist = []
  50         self.zip = None
  51         self.workspace = os.path.normpath(workspace)
  52         self.backupfile = gBACKUPFILE
  53         self.zipfile = os.path.join(self.workspace, gBACKUPFILE)
  54
  55     def _createzip(self):
  56         if self.zip:
  57             return
  58         self.zip = zipfile.ZipFile(self.zipfile, 'w', zipfile.ZIP_DEFLATED)
  59
  60     def _save(self, tmp, path):
  61         if not self._tryhook(path):
  62             return
  63         self.rlist.append(_PathInfo(tmp, path))
  64
  65     def bkrmfile(self, path):
  66         arc = self._tryhook(path)
  67         if arc and os.path.isfile(path):
  68             self._createzip()
  69             self.zip.write(path, arc.encode('utf_8'))
  70             sta = os.stat(path)
  71             oldmode = stat.S_IMODE(sta.st_mode)
  72             self.rlist.append(_PathInfo(_CHMOD, path, oldmode))
  73             self.rlist.append(_PathInfo(_RMFILE, path))
  74         __built_in_remove__(path)
  75
  76     def bkmkfile(self, path, mode, bufsize):
  77         if not os.path.exists(path):
  78             self._save(_MKFILE, path)
  79         return __built_in_open__(path, mode, bufsize)
  80
  81     def bkrmdir(self, path):
  82         if os.path.exists(path):
  83             sta = os.stat(path)
  84             oldmode = stat.S_IMODE(sta.st_mode)
  85             self.rlist.append(_PathInfo(_CHMOD, path, oldmode))
  86             self._save(_RMDIR, path)
  87         __built_in_rmdir__(path)
  88
  89     def bkmkdir(self, path, mode):
  90         if not os.path.exists(path):
  91             self._save(_MKDIR, path)
  92         __built_in_mkdir__(path, mode)
  93
  94     def bkchmod(self, path, mode):
  95         if self._tryhook(path) and os.path.exists(path):
  96             sta = os.stat(path)
  97             oldmode = stat.S_IMODE(sta.st_mode)
  98             self.rlist.append(_PathInfo(_CHMOD, path, oldmode))
  99         __built_in_chmod__(path, mode)
 100
 101     def rollback(self):
 102         if self.zip:
 103             self.zip.close()
 104             self.zip = None
 105         index = len(self.rlist) - 1
 106         while index >= 0:
 107             item = self.rlist[index]
 108             exist = os.path.exists(item.path)
 109             if item.action == _MKFILE and exist:
 110                 #if not os.access(item.path, os.W_OK):
 111                 #    os.chmod(item.path, S_IWUSR)
 112                 __built_in_remove__(item.path)
 113             elif item.action == _RMFILE and not exist:
 114                 if not self.zip:
 115                     self.zip = zipfile.ZipFile(self.zipfile, 'r', zipfile.ZIP_DEFLATED)
 116                 arcname = os.path.normpath(item.path)
 117                 arcname = arcname[len(self.workspace)+1:].encode('utf_8')
 118                 if os.sep != "/" and os.sep in arcname:
 119                     arcname = arcname.replace(os.sep, '/')
 120                 mtime = self.zip.getinfo(arcname).date_time
 121                 content = self.zip.read(arcname)
 122                 filep = __built_in_open__(item.path, "wb")
 123                 filep.write(content)
 124                 filep.close()
 125                 intime = time.mktime(mtime + (0, 0, 0))
 126                 os.utime(item.path, (intime, intime))
 127             elif item.action == _MKDIR and exist:
 128                 while True:
 129                     try:
 130                         __built_in_rmdir__(item.path)
 131                         break
 132                     except IOError:
 133                         # Sleep a short time and try again
 134                         # The anti-virus software may delay the file removal in this directory
 135                         sleep(0.1)
 136             elif item.action == _RMDIR and not exist:
 137                 __built_in_mkdir__(item.path)
 138             elif item.action == _CHMOD and exist:
 139                 try:
 140                     __built_in_chmod__(item.path, item.mode)
 141                 except EnvironmentError:
 142                     pass
 143             index -= 1
 144         self.commit()
 145
 146     def commit(self):
 147         if self.zip:
 148             self.zip.close()
 149             __built_in_remove__(self.zipfile)
 150
 151     # Check if path needs to be hooked
 152     def _tryhook(self, path):
 153         path = os.path.normpath(path)
 154         works = self.workspace if str(self.workspace).endswith(os.sep) else (self.workspace  + os.sep)
 155         if not path.startswith(works):
 156             return ''
 157         for exceptdir in gEXCEPTION_LIST:
 158             full = os.path.join(self.workspace, exceptdir)
 159             if full == path or path.startswith(full + os.sep) or os.path.split(full)[0] == path:
 160                 return ''
 161         return path[len(self.workspace)+1:]
 162
 163 def _hookrm(path):
 164     if GlobalData.gRECOVERMGR:
 165         GlobalData.gRECOVERMGR.bkrmfile(path)
 166     else:
 167         __built_in_remove__(path)
 168
 169 def _hookmkdir(path, mode=0777):
 170     if GlobalData.gRECOVERMGR:
 171         GlobalData.gRECOVERMGR.bkmkdir(path, mode)
 172     else:
 173         __built_in_mkdir__(path, mode)
 174
 175 def _hookrmdir(path):
 176     if GlobalData.gRECOVERMGR:
 177         GlobalData.gRECOVERMGR.bkrmdir(path)
 178     else:
 179         __built_in_rmdir__(path)
 180
 181 def _hookmkfile(path, mode='r', bufsize=-1):
 182     if GlobalData.gRECOVERMGR:
 183         return GlobalData.gRECOVERMGR.bkmkfile(path, mode, bufsize)
 184     return __built_in_open__(path, mode, bufsize)
 185
 186 def _hookchmod(path, mode):
 187     if GlobalData.gRECOVERMGR:
 188         GlobalData.gRECOVERMGR.bkchmod(path, mode)
 189     else:
 190         __built_in_chmod__(path, mode)
 191
 192 def SetRecoverMgr(mgr):
 193     GlobalData.gRECOVERMGR = mgr
 194
 195 os.remove   = _hookrm
 196 os.mkdir    = _hookmkdir
 197 os.rmdir    = _hookrmdir
 198 os.chmod    = _hookchmod
 199 __FileHookOpen__    = _hookmkfile