1# MIT License
2#
3# Copyright The SCons Foundation
4#
5# Permission is hereby granted, free of charge, to any person obtaining
6# a copy of this software and associated documentation files (the
7# "Software"), to deal in the Software without restriction, including
8# without limitation the rights to use, copy, modify, merge, publish,
9# distribute, sublicense, and/or sell copies of the Software, and to
10# permit persons to whom the Software is furnished to do so, subject to
11# the following conditions:
12#
13# The above copyright notice and this permission notice shall be included
14# in all copies or substantial portions of the Software.
15#
16# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
17# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
18# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
19# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
20# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
21# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
22# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
23
24"""
25dblite.py module contributed by Ralf W. Grosse-Kunstleve.
26Extended for Unicode by Steven Knight.
27"""
28
29import os
30import pickle
31import shutil
32import time
33
34from SCons.compat import PICKLE_PROTOCOL
35
36KEEP_ALL_FILES = False
37IGNORE_CORRUPT_DBFILES = False
38
39
40def corruption_warning(filename):
41    """Local warning for corrupt db.
42
43    Used for self-tests. SCons overwrites this with a
44    different warning function in SConsign.py.
45    """
46    print("Warning: Discarding corrupt database:", filename)
47
48DBLITE_SUFFIX = '.dblite'
49TMP_SUFFIX = '.tmp'
50
51
52class dblite:
53    """
54    Squirrel away references to the functions in various modules
55    that we'll use when our __del__() method calls our sync() method
56    during shutdown.  We might get destroyed when Python is in the midst
57    of tearing down the different modules we import in an essentially
58    arbitrary order, and some of the various modules's global attributes
59    may already be wiped out from under us.
60
61    See the discussion at:
62      http://mail.python.org/pipermail/python-bugs-list/2003-March/016877.html
63
64    """
65
66    _open = open
67    _pickle_dump = staticmethod(pickle.dump)
68    _pickle_protocol = PICKLE_PROTOCOL
69
70    try:
71        _os_chown = os.chown
72    except AttributeError:
73        _os_chown = None
74    _os_replace = os.replace
75    _os_chmod = os.chmod
76    _shutil_copyfile = shutil.copyfile
77    _time_time = time.time
78
79    def __init__(self, file_base_name, flag, mode):
80        assert flag in (None, "r", "w", "c", "n")
81        if flag is None:
82            flag = "r"
83
84        base, ext = os.path.splitext(file_base_name)
85        if ext == DBLITE_SUFFIX:
86            # There's already a suffix on the file name, don't add one.
87            self._file_name = file_base_name
88            self._tmp_name = base + TMP_SUFFIX
89        else:
90            self._file_name = file_base_name + DBLITE_SUFFIX
91            self._tmp_name = file_base_name + TMP_SUFFIX
92
93        self._flag = flag
94        self._mode = mode
95        self._dict = {}
96        self._needs_sync = False
97
98        if self._os_chown is not None and (os.geteuid() == 0 or os.getuid() == 0):
99            # running as root; chown back to current owner/group when done
100            try:
101                statinfo = os.stat(self._file_name)
102                self._chown_to = statinfo.st_uid
103                self._chgrp_to = statinfo.st_gid
104            except OSError:
105                # db file doesn't exist yet.
106                # Check os.environ for SUDO_UID, use if set
107                self._chown_to = int(os.environ.get('SUDO_UID', -1))
108                self._chgrp_to = int(os.environ.get('SUDO_GID', -1))
109        else:
110            self._chown_to = -1  # don't chown
111            self._chgrp_to = -1  # don't chgrp
112
113        if self._flag == "n":
114            with self._open(self._file_name, "wb", self._mode):
115                pass  # just make sure it exists
116        else:
117            try:
118                f = self._open(self._file_name, "rb")
119            except IOError as e:
120                if self._flag != "c":
121                    raise e
122                with self._open(self._file_name, "wb", self._mode):
123                    pass  # just make sure it exists
124            else:
125                p = f.read()
126                f.close()
127                if len(p) > 0:
128                    try:
129                        self._dict = pickle.loads(p, encoding='bytes')
130                    except (pickle.UnpicklingError, EOFError, KeyError):
131                        # Note how we catch KeyErrors too here, which might happen
132                        # when we don't have cPickle available (default pickle
133                        # throws it).
134                        if IGNORE_CORRUPT_DBFILES:
135                            corruption_warning(self._file_name)
136                        else:
137                            raise
138
139    def close(self):
140        if self._needs_sync:
141            self.sync()
142
143    def __del__(self):
144        self.close()
145
146    def sync(self):
147        self._check_writable()
148        with self._open(self._tmp_name, "wb", self._mode) as f:
149            self._pickle_dump(self._dict, f, self._pickle_protocol)
150
151        try:
152            self._os_replace(self._tmp_name, self._file_name)
153        except PermissionError:
154            # If we couldn't replace due to perms, try to change and retry.
155            # This is mainly for Windows - on POSIX the file permissions
156            # don't matter, the os.replace would have worked anyway.
157            # We're giving up if the retry fails, just let the Python
158            # exception abort us.
159            try:
160                self._os_chmod(self._file_name, 0o777)
161            except PermissionError:
162                pass
163            self._os_replace(self._tmp_name, self._file_name)
164
165        if self._os_chown is not None and self._chown_to > 0:  # don't chown to root or -1
166            try:
167                self._os_chown(self._file_name, self._chown_to, self._chgrp_to)
168            except OSError:
169                pass
170
171        self._needs_sync = False
172        if KEEP_ALL_FILES:
173            self._shutil_copyfile(
174                self._file_name,
175                self._file_name + "_" + str(int(self._time_time()))
176            )
177
178    def _check_writable(self):
179        if self._flag == "r":
180            raise IOError("Read-only database: %s" % self._file_name)
181
182    def __getitem__(self, key):
183        return self._dict[key]
184
185    def __setitem__(self, key, value):
186        self._check_writable()
187
188        if not isinstance(key, str):
189            raise TypeError("key `%s' must be a string but is %s" % (key, type(key)))
190
191        if not isinstance(value, bytes):
192            raise TypeError("value `%s' must be a bytes but is %s" % (value, type(value)))
193
194        self._dict[key] = value
195        self._needs_sync = True
196
197    def keys(self):
198        return list(self._dict.keys())
199
200    def __contains__(self, key):
201        return key in self._dict
202
203    def __iter__(self):
204        return iter(self._dict)
205
206    def __len__(self):
207        return len(self._dict)
208
209
210def open(file, flag=None, mode=0o666):
211    return dblite(file, flag, mode)
212
213
214def _exercise():
215    db = open("tmp", "n")
216    assert len(db) == 0
217    db["foo"] = b"bar"
218    assert db["foo"] == b"bar"
219    db.sync()
220
221    db = open("tmp", "c")
222    assert len(db) == 1, len(db)
223    assert db["foo"] == b"bar"
224    db["bar"] = b"foo"
225    assert db["bar"] == b"foo"
226    db.sync()
227
228    db = open("tmp", "r")
229    assert len(db) == 2, len(db)
230    assert db["foo"] == b"bar"
231    assert db["bar"] == b"foo"
232    try:
233        db.sync()
234    except IOError as e:
235        assert str(e) == "Read-only database: tmp.dblite"
236    else:
237        raise RuntimeError("IOError expected.")
238    db = open("tmp", "w")
239    assert len(db) == 2, len(db)
240    db["ping"] = b"pong"
241    db.sync()
242
243    try:
244        db[(1, 2)] = "tuple"
245    except TypeError as e:
246        assert str(e) == "key `(1, 2)' must be a string but is <class 'tuple'>", str(e)
247    else:
248        raise RuntimeError("TypeError exception expected")
249
250    try:
251        db["list"] = [1, 2]
252    except TypeError as e:
253        assert str(e) == "value `[1, 2]' must be a bytes but is <class 'list'>", str(e)
254    else:
255        raise RuntimeError("TypeError exception expected")
256
257    db = open("tmp", "r")
258    assert len(db) == 3, len(db)
259
260    db = open("tmp", "n")
261    assert len(db) == 0, len(db)
262    dblite._open("tmp.dblite", "w")
263
264    db = open("tmp", "r")
265    dblite._open("tmp.dblite", "w").write("x")
266    try:
267        db = open("tmp", "r")
268    except pickle.UnpicklingError:
269        pass
270    else:
271        raise RuntimeError("pickle exception expected.")
272
273    global IGNORE_CORRUPT_DBFILES
274    IGNORE_CORRUPT_DBFILES = True
275    db = open("tmp", "r")
276    assert len(db) == 0, len(db)
277    os.unlink("tmp.dblite")
278    try:
279        db = open("tmp", "w")
280    except IOError as e:
281        assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e)
282    else:
283        raise RuntimeError("IOError expected.")
284
285    print("Completed _exercise()")
286
287
288if __name__ == "__main__":
289    _exercise()
290
291# Local Variables:
292# tab-width:4
293# indent-tabs-mode:nil
294# End:
295# vim: set expandtab tabstop=4 shiftwidth=4:
296