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