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'''
16File hook
17'''
18
19import os
20import stat
21import time
22import zipfile
23from time import sleep
24from 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
38gBACKUPFILE = 'file.backup'
39gEXCEPTION_LIST = ['Conf'+os.sep+'DistributionPackageDatabase.db', '.tmp', gBACKUPFILE]
40
41class _PathInfo:
42    def __init__(self, action, path, mode=-1):
43        self.action = action
44        self.path = path
45        self.mode = mode
46
47class 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
163def _hookrm(path):
164    if GlobalData.gRECOVERMGR:
165        GlobalData.gRECOVERMGR.bkrmfile(path)
166    else:
167        __built_in_remove__(path)
168
169def _hookmkdir(path, mode=0777):
170    if GlobalData.gRECOVERMGR:
171        GlobalData.gRECOVERMGR.bkmkdir(path, mode)
172    else:
173        __built_in_mkdir__(path, mode)
174
175def _hookrmdir(path):
176    if GlobalData.gRECOVERMGR:
177        GlobalData.gRECOVERMGR.bkrmdir(path)
178    else:
179        __built_in_rmdir__(path)
180
181def _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
186def _hookchmod(path, mode):
187    if GlobalData.gRECOVERMGR:
188        GlobalData.gRECOVERMGR.bkchmod(path, mode)
189    else:
190        __built_in_chmod__(path, mode)
191
192def SetRecoverMgr(mgr):
193    GlobalData.gRECOVERMGR = mgr
194
195os.remove   = _hookrm
196os.mkdir    = _hookmkdir
197os.rmdir    = _hookrmdir
198os.chmod    = _hookchmod
199__FileHookOpen__    = _hookmkfile
200