1""" 2A module written originally by Armin Ronacher to manage file transfers in an 3atomic way 4""" 5import errno 6import os 7import random 8import shutil 9import sys 10import tempfile 11import time 12 13import salt.utils.win_dacl 14 15CAN_RENAME_OPEN_FILE = False 16if os.name == "nt": # pragma: no cover 17 _rename = lambda src, dst: False # pylint: disable=C0103 18 _rename_atomic = lambda src, dst: False # pylint: disable=C0103 19 20 try: 21 import ctypes 22 23 _MOVEFILE_REPLACE_EXISTING = 0x1 24 _MOVEFILE_WRITE_THROUGH = 0x8 25 _MoveFileEx = ctypes.windll.kernel32.MoveFileExW # pylint: disable=C0103 26 27 def _rename(src, dst): # pylint: disable=E0102 28 if not isinstance(src, str): 29 src = str(src, sys.getfilesystemencoding()) 30 if not isinstance(dst, str): 31 dst = str(dst, sys.getfilesystemencoding()) 32 if _rename_atomic(src, dst): 33 return True 34 retry = 0 35 rval = False 36 while not rval and retry < 100: 37 rval = _MoveFileEx( 38 src, dst, _MOVEFILE_REPLACE_EXISTING | _MOVEFILE_WRITE_THROUGH 39 ) 40 if not rval: 41 time.sleep(0.001) 42 retry += 1 43 return rval 44 45 # new in Vista and Windows Server 2008 46 # pylint: disable=C0103 47 _CreateTransaction = ctypes.windll.ktmw32.CreateTransaction 48 _CommitTransaction = ctypes.windll.ktmw32.CommitTransaction 49 _MoveFileTransacted = ctypes.windll.kernel32.MoveFileTransactedW 50 _CloseHandle = ctypes.windll.kernel32.CloseHandle 51 # pylint: enable=C0103 52 CAN_RENAME_OPEN_FILE = True 53 54 def _rename_atomic(src, dst): # pylint: disable=E0102 55 tra = _CreateTransaction(None, 0, 0, 0, 0, 1000, "Atomic rename") 56 if tra == -1: 57 return False 58 try: 59 retry = 0 60 rval = False 61 while not rval and retry < 100: 62 rval = _MoveFileTransacted( 63 src, 64 dst, 65 None, 66 None, 67 _MOVEFILE_REPLACE_EXISTING | _MOVEFILE_WRITE_THROUGH, 68 tra, 69 ) 70 if rval: 71 rval = _CommitTransaction(tra) 72 break 73 else: 74 time.sleep(0.001) 75 retry += 1 76 return rval 77 finally: 78 _CloseHandle(tra) 79 80 except Exception: # pylint: disable=broad-except 81 pass 82 83 def atomic_rename(src, dst): 84 # Try atomic or pseudo-atomic rename 85 if _rename(src, dst): 86 return 87 # Fall back to "move away and replace" 88 try: 89 os.rename(src, dst) 90 except OSError as err: 91 if err.errno != errno.EEXIST: 92 raise 93 old = "{}-{:08x}".format(dst, random.randint(0, sys.maxint)) 94 os.rename(dst, old) 95 os.rename(src, dst) 96 try: 97 os.unlink(old) 98 except Exception: # pylint: disable=broad-except 99 pass 100 101 102else: 103 atomic_rename = os.rename # pylint: disable=C0103 104 CAN_RENAME_OPEN_FILE = True 105 106 107class _AtomicWFile: 108 """ 109 Helper class for :func:`atomic_open`. 110 """ 111 112 def __init__(self, fhanle, tmp_filename, filename): 113 self._fh = fhanle 114 self._tmp_filename = tmp_filename 115 self._filename = filename 116 117 def __getattr__(self, attr): 118 return getattr(self._fh, attr) 119 120 def __enter__(self): 121 return self 122 123 def close(self): 124 if self._fh.closed: 125 return 126 self._fh.close() 127 if os.path.isfile(self._filename): 128 if salt.utils.win_dacl.HAS_WIN32: 129 salt.utils.win_dacl.copy_security( 130 source=self._filename, target=self._tmp_filename 131 ) 132 else: 133 shutil.copymode(self._filename, self._tmp_filename) 134 st = os.stat(self._filename) 135 os.chown(self._tmp_filename, st.st_uid, st.st_gid) 136 atomic_rename(self._tmp_filename, self._filename) 137 138 def __exit__(self, exc_type, exc_value, traceback): 139 if exc_type is None: 140 self.close() 141 else: 142 self._fh.close() 143 try: 144 os.remove(self._tmp_filename) 145 except OSError: 146 pass 147 148 def __repr__(self): 149 return "<{} {}{}, mode {}>".format( 150 self.__class__.__name__, 151 self._fh.closed and "closed " or "", 152 self._filename, 153 self._fh.mode, 154 ) 155 156 157def atomic_open(filename, mode="w"): 158 """ 159 Works like a regular `open()` but writes updates into a temporary 160 file instead of the given file and moves it over when the file is 161 closed. The file returned behaves as if it was a regular Python 162 """ 163 if mode in ("r", "rb", "r+", "rb+", "a", "ab"): 164 raise TypeError("Read or append modes don't work with atomic_open") 165 kwargs = { 166 "prefix": ".___atomic_write", 167 "dir": os.path.dirname(filename), 168 "delete": False, 169 } 170 if "b" not in mode: 171 kwargs["newline"] = "" 172 ntf = tempfile.NamedTemporaryFile(mode, **kwargs) 173 return _AtomicWFile(ntf, ntf.name, filename) 174