OLD | NEW |
| (Empty) |
1 # -*- test-case-name: twisted.test.test_sob -*- | |
2 # Copyright (c) 2001-2004 Twisted Matrix Laboratories. | |
3 # See LICENSE for details. | |
4 | |
5 # | |
6 """ | |
7 Save and load Small OBjects to and from files, using various formats. | |
8 | |
9 Maintainer: U{Moshe Zadka<mailto:moshez@twistedmatrix.com>} | |
10 """ | |
11 | |
12 import hashlib | |
13 import os, sys | |
14 try: | |
15 import cPickle as pickle | |
16 except ImportError: | |
17 import pickle | |
18 try: | |
19 import cStringIO as StringIO | |
20 except ImportError: | |
21 import StringIO | |
22 from twisted.python import log, runtime | |
23 from twisted.persisted import styles | |
24 from zope.interface import implements, Interface | |
25 | |
26 # Note: | |
27 # These encrypt/decrypt functions only work for data formats | |
28 # which are immune to having spaces tucked at the end. | |
29 # All data formats which persist saves hold that condition. | |
30 def _encrypt(passphrase, data): | |
31 from Crypto.Cipher import AES as cipher | |
32 leftover = len(data) % cipher.block_size | |
33 if leftover: | |
34 data += ' '*(cipher.block_size - leftover) | |
35 return cipher.new(hashlib.md5(passphrase).digest()[:16]).encrypt(data) | |
36 | |
37 def _decrypt(passphrase, data): | |
38 from Crypto.Cipher import AES | |
39 return AES.new(hashlib.md5(passphrase).digest()[:16]).decrypt(data) | |
40 | |
41 | |
42 class IPersistable(Interface): | |
43 | |
44 """An object which can be saved in several formats to a file""" | |
45 | |
46 def setStyle(style): | |
47 """Set desired format. | |
48 | |
49 @type style: string (one of 'pickle', 'source' or 'xml') | |
50 """ | |
51 | |
52 def save(tag=None, filename=None, passphrase=None): | |
53 """Save object to file. | |
54 | |
55 @type tag: string | |
56 @type filename: string | |
57 @type passphrase: string | |
58 """ | |
59 | |
60 | |
61 class Persistent: | |
62 | |
63 implements(IPersistable) | |
64 | |
65 style = "pickle" | |
66 | |
67 def __init__(self, original, name): | |
68 self.original = original | |
69 self.name = name | |
70 | |
71 def setStyle(self, style): | |
72 """Set desired format. | |
73 | |
74 @type style: string (one of 'pickle', 'source' or 'xml') | |
75 """ | |
76 self.style = style | |
77 | |
78 def _getFilename(self, filename, ext, tag): | |
79 if filename: | |
80 finalname = filename | |
81 filename = finalname + "-2" | |
82 elif tag: | |
83 filename = "%s-%s-2.%s" % (self.name, tag, ext) | |
84 finalname = "%s-%s.%s" % (self.name, tag, ext) | |
85 else: | |
86 filename = "%s-2.%s" % (self.name, ext) | |
87 finalname = "%s.%s" % (self.name, ext) | |
88 return finalname, filename | |
89 | |
90 def _saveTemp(self, filename, passphrase, dumpFunc): | |
91 f = open(filename, 'wb') | |
92 if passphrase is None: | |
93 dumpFunc(self.original, f) | |
94 else: | |
95 s = StringIO.StringIO() | |
96 dumpFunc(self.original, s) | |
97 f.write(_encrypt(passphrase, s.getvalue())) | |
98 f.close() | |
99 | |
100 def _getStyle(self): | |
101 if self.style == "xml": | |
102 from twisted.persisted.marmalade import jellyToXML as dumpFunc | |
103 ext = "tax" | |
104 elif self.style == "source": | |
105 from twisted.persisted.aot import jellyToSource as dumpFunc | |
106 ext = "tas" | |
107 else: | |
108 def dumpFunc(obj, file): | |
109 pickle.dump(obj, file, 2) | |
110 ext = "tap" | |
111 return ext, dumpFunc | |
112 | |
113 def save(self, tag=None, filename=None, passphrase=None): | |
114 """Save object to file. | |
115 | |
116 @type tag: string | |
117 @type filename: string | |
118 @type passphrase: string | |
119 """ | |
120 ext, dumpFunc = self._getStyle() | |
121 if passphrase: | |
122 ext = 'e' + ext | |
123 finalname, filename = self._getFilename(filename, ext, tag) | |
124 log.msg("Saving "+self.name+" application to "+finalname+"...") | |
125 self._saveTemp(filename, passphrase, dumpFunc) | |
126 if runtime.platformType == "win32" and os.path.isfile(finalname): | |
127 os.remove(finalname) | |
128 os.rename(filename, finalname) | |
129 log.msg("Saved.") | |
130 | |
131 # "Persistant" has been present since 1.0.7, so retain it for compatibility | |
132 Persistant = Persistent | |
133 | |
134 class _EverythingEphemeral(styles.Ephemeral): | |
135 | |
136 initRun = 0 | |
137 | |
138 def __init__(self, mainMod): | |
139 """ | |
140 @param mainMod: The '__main__' module that this class will proxy. | |
141 """ | |
142 self.mainMod = mainMod | |
143 | |
144 def __getattr__(self, key): | |
145 try: | |
146 return getattr(self.mainMod, key) | |
147 except AttributeError: | |
148 if self.initRun: | |
149 raise | |
150 else: | |
151 log.msg("Warning! Loading from __main__: %s" % key) | |
152 return styles.Ephemeral() | |
153 | |
154 | |
155 def load(filename, style, passphrase=None): | |
156 """Load an object from a file. | |
157 | |
158 Deserialize an object from a file. The file can be encrypted. | |
159 | |
160 @param filename: string | |
161 @param style: string (one of 'source', 'xml' or 'pickle') | |
162 @param passphrase: string | |
163 """ | |
164 mode = 'r' | |
165 if style=='source': | |
166 from twisted.persisted.aot import unjellyFromSource as _load | |
167 elif style=='xml': | |
168 from twisted.persisted.marmalade import unjellyFromXML as _load | |
169 else: | |
170 _load, mode = pickle.load, 'rb' | |
171 if passphrase: | |
172 fp = StringIO.StringIO(_decrypt(passphrase, | |
173 open(filename, 'rb').read())) | |
174 else: | |
175 fp = open(filename, mode) | |
176 ee = _EverythingEphemeral(sys.modules['__main__']) | |
177 sys.modules['__main__'] = ee | |
178 ee.initRun = 1 | |
179 try: | |
180 value = _load(fp) | |
181 finally: | |
182 # restore __main__ if an exception is raised. | |
183 sys.modules['__main__'] = ee.mainMod | |
184 | |
185 styles.doUpgrade() | |
186 ee.initRun = 0 | |
187 persistable = IPersistable(value, None) | |
188 if persistable is not None: | |
189 persistable.setStyle(style) | |
190 return value | |
191 | |
192 | |
193 def loadValueFromFile(filename, variable, passphrase=None): | |
194 """Load the value of a variable in a Python file. | |
195 | |
196 Run the contents of the file, after decrypting if C{passphrase} is | |
197 given, in a namespace and return the result of the variable | |
198 named C{variable}. | |
199 | |
200 @param filename: string | |
201 @param variable: string | |
202 @param passphrase: string | |
203 """ | |
204 if passphrase: | |
205 mode = 'rb' | |
206 else: | |
207 mode = 'r' | |
208 fileObj = open(filename, mode) | |
209 d = {'__file__': filename} | |
210 if passphrase: | |
211 data = fileObj.read() | |
212 data = _decrypt(passphrase, data) | |
213 exec data in d, d | |
214 else: | |
215 exec fileObj in d, d | |
216 value = d[variable] | |
217 return value | |
218 | |
219 def guessType(filename): | |
220 ext = os.path.splitext(filename)[1] | |
221 return { | |
222 '.tac': 'python', | |
223 '.etac': 'python', | |
224 '.py': 'python', | |
225 '.tap': 'pickle', | |
226 '.etap': 'pickle', | |
227 '.tas': 'source', | |
228 '.etas': 'source', | |
229 '.tax': 'xml', | |
230 '.etax': 'xml' | |
231 }[ext] | |
232 | |
233 __all__ = ['loadValueFromFile', 'load', 'Persistent', 'Persistant', | |
234 'IPersistable', 'guessType'] | |
OLD | NEW |