1# -*- test-case-name: twisted.test.test_dirdbm -*-
2#
3# Copyright (c) Twisted Matrix Laboratories.
4# See LICENSE for details.
5
6
7"""
8DBM-style interface to a directory.
9
10Each key is stored as a single file.  This is not expected to be very fast or
11efficient, but it's good for easy debugging.
12
13DirDBMs are *not* thread-safe, they should only be accessed by one thread at
14a time.
15
16No files should be placed in the working directory of a DirDBM save those
17created by the DirDBM itself!
18
19Maintainer: Itamar Shtull-Trauring
20"""
21
22
23import base64
24import glob
25import os
26import pickle
27
28from twisted.python.filepath import FilePath
29
30try:
31    _open  # type: ignore[has-type]
32except NameError:
33    _open = open
34
35
36class DirDBM:
37    """
38    A directory with a DBM interface.
39
40    This class presents a hash-like interface to a directory of small,
41    flat files. It can only use strings as keys or values.
42    """
43
44    def __init__(self, name):
45        """
46        @type name: str
47        @param name: Base path to use for the directory storage.
48        """
49        self.dname = os.path.abspath(name)
50        self._dnamePath = FilePath(name)
51        if not self._dnamePath.isdir():
52            self._dnamePath.createDirectory()
53        else:
54            # Run recovery, in case we crashed. we delete all files ending
55            # with ".new". Then we find all files who end with ".rpl". If a
56            # corresponding file exists without ".rpl", we assume the write
57            # failed and delete the ".rpl" file. If only a ".rpl" exist we
58            # assume the program crashed right after deleting the old entry
59            # but before renaming the replacement entry.
60            #
61            # NOTE: '.' is NOT in the base64 alphabet!
62            for f in glob.glob(self._dnamePath.child("*.new").path):
63                os.remove(f)
64            replacements = glob.glob(self._dnamePath.child("*.rpl").path)
65            for f in replacements:
66                old = f[:-4]
67                if os.path.exists(old):
68                    os.remove(f)
69                else:
70                    os.rename(f, old)
71
72    def _encode(self, k):
73        """
74        Encode a key so it can be used as a filename.
75        """
76        # NOTE: '_' is NOT in the base64 alphabet!
77        return base64.encodebytes(k).replace(b"\n", b"_").replace(b"/", b"-")
78
79    def _decode(self, k):
80        """
81        Decode a filename to get the key.
82        """
83        return base64.decodebytes(k.replace(b"_", b"\n").replace(b"-", b"/"))
84
85    def _readFile(self, path):
86        """
87        Read in the contents of a file.
88
89        Override in subclasses to e.g. provide transparently encrypted dirdbm.
90        """
91        with _open(path.path, "rb") as f:
92            s = f.read()
93        return s
94
95    def _writeFile(self, path, data):
96        """
97        Write data to a file.
98
99        Override in subclasses to e.g. provide transparently encrypted dirdbm.
100        """
101        with _open(path.path, "wb") as f:
102            f.write(data)
103            f.flush()
104
105    def __len__(self):
106        """
107        @return: The number of key/value pairs in this Shelf
108        """
109        return len(self._dnamePath.listdir())
110
111    def __setitem__(self, k, v):
112        """
113        C{dirdbm[k] = v}
114        Create or modify a textfile in this directory
115
116        @type k: bytes
117        @param k: key to set
118
119        @type v: bytes
120        @param v: value to associate with C{k}
121        """
122        if not type(k) == bytes:
123            raise TypeError("DirDBM key must be bytes")
124        if not type(v) == bytes:
125            raise TypeError("DirDBM value must be bytes")
126        k = self._encode(k)
127
128        # We create a new file with extension .new, write the data to it, and
129        # if the write succeeds delete the old file and rename the new one.
130        old = self._dnamePath.child(k)
131        if old.exists():
132            new = old.siblingExtension(".rpl")  # Replacement entry
133        else:
134            new = old.siblingExtension(".new")  # New entry
135        try:
136            self._writeFile(new, v)
137        except BaseException:
138            new.remove()
139            raise
140        else:
141            if old.exists():
142                old.remove()
143            new.moveTo(old)
144
145    def __getitem__(self, k):
146        """
147        C{dirdbm[k]}
148        Get the contents of a file in this directory as a string.
149
150        @type k: bytes
151        @param k: key to lookup
152
153        @return: The value associated with C{k}
154        @raise KeyError: Raised when there is no such key
155        """
156        if not type(k) == bytes:
157            raise TypeError("DirDBM key must be bytes")
158        path = self._dnamePath.child(self._encode(k))
159        try:
160            return self._readFile(path)
161        except (OSError):
162            raise KeyError(k)
163
164    def __delitem__(self, k):
165        """
166        C{del dirdbm[foo]}
167        Delete a file in this directory.
168
169        @type k: bytes
170        @param k: key to delete
171
172        @raise KeyError: Raised when there is no such key
173        """
174        if not type(k) == bytes:
175            raise TypeError("DirDBM key must be bytes")
176        k = self._encode(k)
177        try:
178            self._dnamePath.child(k).remove()
179        except (OSError):
180            raise KeyError(self._decode(k))
181
182    def keys(self):
183        """
184        @return: a L{list} of filenames (keys).
185        """
186        return list(map(self._decode, self._dnamePath.asBytesMode().listdir()))
187
188    def values(self):
189        """
190        @return: a L{list} of file-contents (values).
191        """
192        vals = []
193        keys = self.keys()
194        for key in keys:
195            vals.append(self[key])
196        return vals
197
198    def items(self):
199        """
200        @return: a L{list} of 2-tuples containing key/value pairs.
201        """
202        items = []
203        keys = self.keys()
204        for key in keys:
205            items.append((key, self[key]))
206        return items
207
208    def has_key(self, key):
209        """
210        @type key: bytes
211        @param key: The key to test
212
213        @return: A true value if this dirdbm has the specified key, a false
214        value otherwise.
215        """
216        if not type(key) == bytes:
217            raise TypeError("DirDBM key must be bytes")
218        key = self._encode(key)
219        return self._dnamePath.child(key).isfile()
220
221    def setdefault(self, key, value):
222        """
223        @type key: bytes
224        @param key: The key to lookup
225
226        @param value: The value to associate with key if key is not already
227        associated with a value.
228        """
229        if key not in self:
230            self[key] = value
231            return value
232        return self[key]
233
234    def get(self, key, default=None):
235        """
236        @type key: bytes
237        @param key: The key to lookup
238
239        @param default: The value to return if the given key does not exist
240
241        @return: The value associated with C{key} or C{default} if not
242        L{DirDBM.has_key(key)}
243        """
244        if key in self:
245            return self[key]
246        else:
247            return default
248
249    def __contains__(self, key):
250        """
251        @see: L{DirDBM.has_key}
252        """
253        return self.has_key(key)
254
255    def update(self, dict):
256        """
257        Add all the key/value pairs in L{dict} to this dirdbm.  Any conflicting
258        keys will be overwritten with the values from L{dict}.
259
260        @type dict: mapping
261        @param dict: A mapping of key/value pairs to add to this dirdbm.
262        """
263        for key, val in dict.items():
264            self[key] = val
265
266    def copyTo(self, path):
267        """
268        Copy the contents of this dirdbm to the dirdbm at C{path}.
269
270        @type path: L{str}
271        @param path: The path of the dirdbm to copy to.  If a dirdbm
272        exists at the destination path, it is cleared first.
273
274        @rtype: C{DirDBM}
275        @return: The dirdbm this dirdbm was copied to.
276        """
277        path = FilePath(path)
278        assert path != self._dnamePath
279
280        d = self.__class__(path.path)
281        d.clear()
282        for k in self.keys():
283            d[k] = self[k]
284        return d
285
286    def clear(self):
287        """
288        Delete all key/value pairs in this dirdbm.
289        """
290        for k in self.keys():
291            del self[k]
292
293    def close(self):
294        """
295        Close this dbm: no-op, for dbm-style interface compliance.
296        """
297
298    def getModificationTime(self, key):
299        """
300        Returns modification time of an entry.
301
302        @return: Last modification date (seconds since epoch) of entry C{key}
303        @raise KeyError: Raised when there is no such key
304        """
305        if not type(key) == bytes:
306            raise TypeError("DirDBM key must be bytes")
307        path = self._dnamePath.child(self._encode(key))
308        if path.isfile():
309            return path.getModificationTime()
310        else:
311            raise KeyError(key)
312
313
314class Shelf(DirDBM):
315    """
316    A directory with a DBM shelf interface.
317
318    This class presents a hash-like interface to a directory of small,
319    flat files. Keys must be strings, but values can be any given object.
320    """
321
322    def __setitem__(self, k, v):
323        """
324        C{shelf[foo] = bar}
325        Create or modify a textfile in this directory.
326
327        @type k: str
328        @param k: The key to set
329
330        @param v: The value to associate with C{key}
331        """
332        v = pickle.dumps(v)
333        DirDBM.__setitem__(self, k, v)
334
335    def __getitem__(self, k):
336        """
337        C{dirdbm[foo]}
338        Get and unpickle the contents of a file in this directory.
339
340        @type k: bytes
341        @param k: The key to lookup
342
343        @return: The value associated with the given key
344        @raise KeyError: Raised if the given key does not exist
345        """
346        return pickle.loads(DirDBM.__getitem__(self, k))
347
348
349def open(file, flag=None, mode=None):
350    """
351    This is for 'anydbm' compatibility.
352
353    @param file: The parameter to pass to the DirDBM constructor.
354
355    @param flag: ignored
356    @param mode: ignored
357    """
358    return DirDBM(file)
359
360
361__all__ = ["open", "DirDBM", "Shelf"]
362