[mirror_edk2.git] / BaseTools / Source / Python / UPT / Core / FileHook.py

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