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