1# Copyright (c) gocept gmbh & co. kg
2# See also LICENSE.txt
3
4"""Persistent dict to remember state between invocations.
5
6Cookies are used to remember file positions, counters and the like
7between plugin invocations. It is not intended for substantial amounts
8of data. Cookies are serialized into JSON and saved to a state file. We
9prefer a plain text format to allow administrators to inspect and edit
10its content. See :class:`~nagiosplugin.logtail.LogTail` for an
11application of cookies to get only new lines of a continuously growing
12file.
13
14Cookies are locked exclusively so that at most one process at a time has
15access to it. Changes to the dict are not reflected in the file until
16:meth:`Cookie.commit` is called. It is recommended to use Cookie as
17context manager to get it opened and committed automatically.
18"""
19
20from .compat import UserDict, TemporaryFile
21from .platform import flock_exclusive
22import codecs
23import json
24import os
25import tempfile
26
27
28class Cookie(UserDict, object):
29
30    def __init__(self, statefile=None):
31        """Creates a persistent dict to keep state.
32
33        After creation, a cookie behaves like a normal dict.
34
35        :param statefile: file name to save the dict's contents
36
37        .. note:: If `statefile` is empty or None, the Cookie will be
38           oblivous, i.e., it will forget its contents on garbage
39           collection. This makes it possible to explicitely throw away
40           state between plugin runs (for example by a command line
41           argument).
42        """
43        super(Cookie, self).__init__()
44        self.path = statefile
45        self.fobj = None
46
47    def __enter__(self):
48        """Allows Cookie to be used as context manager.
49
50        Opens the file and passes a dict-like object into the
51        subordinate context. See :meth:`open` for details about opening
52        semantics. When the context is left in the regular way (no
53        exception raised), the cookie is :meth:`commit`\ ted to disk.
54
55        :yields: open cookie
56        """
57        self.open()
58        return self
59
60    def __exit__(self, exc_type, exc_val, exc_tb):
61        if not exc_type:
62            self.commit()
63        self.close()
64
65    def open(self):
66        """Reads/creates the state file and initializes the dict.
67
68        If the state file does not exist, it is touched into existence.
69        An exclusive lock is acquired to ensure serialized access. If
70        :meth:`open` fails to parse file contents, it truncates
71        the file before raising an exception. This guarantees that
72        plugins will not fail repeatedly when their state files get
73        damaged.
74
75        :returns: Cookie object (self)
76        :raises ValueError: if the state file is corrupted or does not
77            deserialize into a dict
78        """
79        self.fobj = self._create_fobj()
80        flock_exclusive(self.fobj)
81        if os.fstat(self.fobj.fileno()).st_size:
82            try:
83                self.data = self._load()
84            except ValueError:
85                self.fobj.truncate(0)
86                raise
87        return self
88
89    def _create_fobj(self):
90        if not self.path:
91            return TemporaryFile('w+', encoding='ascii',
92                                 prefix='oblivious_cookie_')
93        # mode='a+' has problems with mixed R/W operation on Mac OS X
94        try:
95            return codecs.open(self.path, 'r+', encoding='ascii')
96        except IOError:
97            return codecs.open(self.path, 'w+', encoding='ascii')
98
99    def _load(self):
100        self.fobj.seek(0)
101        data = json.load(self.fobj)
102        if not isinstance(data, dict):
103            raise ValueError('format error: cookie does not contain dict',
104                             self.path, data)
105        return data
106
107    def close(self):
108        """Closes a cookie and its underlying state file.
109
110        This method has no effect if the cookie is already closed.
111        Once the cookie is closed, any operation (like :meth:`commit`)
112        will raise an exception.
113        """
114        if not self.fobj:
115            return
116        self.fobj.close()
117        self.fobj = None
118
119    def commit(self):
120        """Persists the cookie's dict items in the state file.
121
122        The cookies content is serialized as JSON string and saved to
123        the state file. The buffers are flushed to ensure that the new
124        content is saved in a durable way.
125        """
126        if not self.fobj:
127            raise IOError('cannot commit closed cookie', self.path)
128        self.fobj.seek(0)
129        self.fobj.truncate()
130        json.dump(self.data, self.fobj)
131        self.fobj.write('\n')
132        self.fobj.flush()
133        os.fsync(self.fobj)
134