1
# -*- test-case-name: twisted.test.test_sob -*-
2
# Copyright (c) 2001-2008 Twisted Matrix Laboratories.
3
# See LICENSE for details.
7
Save and load Small OBjects to and from files, using various formats.
9
Maintainer: Moshe Zadka
14
import cPickle as pickle
18
import cStringIO as StringIO
21
from twisted.python import log, runtime
22
from twisted.python.hashlib import md5
23
from twisted.persisted import styles
24
from zope.interface import implements, Interface
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
34
data += ' '*(cipher.block_size - leftover)
35
return cipher.new(md5(passphrase).digest()[:16]).encrypt(data)
37
def _decrypt(passphrase, data):
38
from Crypto.Cipher import AES
39
return AES.new(md5(passphrase).digest()[:16]).decrypt(data)
42
class IPersistable(Interface):
44
"""An object which can be saved in several formats to a file"""
47
"""Set desired format.
49
@type style: string (one of 'pickle' or 'source')
52
def save(tag=None, filename=None, passphrase=None):
53
"""Save object to file.
56
@type filename: string
57
@type passphrase: string
63
implements(IPersistable)
67
def __init__(self, original, name):
68
self.original = original
71
def setStyle(self, style):
72
"""Set desired format.
74
@type style: string (one of 'pickle' or 'source')
78
def _getFilename(self, filename, ext, tag):
81
filename = finalname + "-2"
83
filename = "%s-%s-2.%s" % (self.name, tag, ext)
84
finalname = "%s-%s.%s" % (self.name, tag, ext)
86
filename = "%s-2.%s" % (self.name, ext)
87
finalname = "%s.%s" % (self.name, ext)
88
return finalname, filename
90
def _saveTemp(self, filename, passphrase, dumpFunc):
91
f = open(filename, 'wb')
92
if passphrase is None:
93
dumpFunc(self.original, f)
95
s = StringIO.StringIO()
96
dumpFunc(self.original, s)
97
f.write(_encrypt(passphrase, s.getvalue()))
101
if self.style == "source":
102
from twisted.persisted.aot import jellyToSource as dumpFunc
105
def dumpFunc(obj, file):
106
pickle.dump(obj, file, 2)
110
def save(self, tag=None, filename=None, passphrase=None):
111
"""Save object to file.
114
@type filename: string
115
@type passphrase: string
117
ext, dumpFunc = self._getStyle()
120
finalname, filename = self._getFilename(filename, ext, tag)
121
log.msg("Saving "+self.name+" application to "+finalname+"...")
122
self._saveTemp(filename, passphrase, dumpFunc)
123
if runtime.platformType == "win32" and os.path.isfile(finalname):
125
os.rename(filename, finalname)
128
# "Persistant" has been present since 1.0.7, so retain it for compatibility
129
Persistant = Persistent
131
class _EverythingEphemeral(styles.Ephemeral):
135
def __init__(self, mainMod):
137
@param mainMod: The '__main__' module that this class will proxy.
139
self.mainMod = mainMod
141
def __getattr__(self, key):
143
return getattr(self.mainMod, key)
144
except AttributeError:
148
log.msg("Warning! Loading from __main__: %s" % key)
149
return styles.Ephemeral()
152
def load(filename, style, passphrase=None):
153
"""Load an object from a file.
155
Deserialize an object from a file. The file can be encrypted.
157
@param filename: string
158
@param style: string (one of 'pickle' or 'source')
159
@param passphrase: string
163
from twisted.persisted.aot import unjellyFromSource as _load
165
_load, mode = pickle.load, 'rb'
167
fp = StringIO.StringIO(_decrypt(passphrase,
168
open(filename, 'rb').read()))
170
fp = open(filename, mode)
171
ee = _EverythingEphemeral(sys.modules['__main__'])
172
sys.modules['__main__'] = ee
177
# restore __main__ if an exception is raised.
178
sys.modules['__main__'] = ee.mainMod
182
persistable = IPersistable(value, None)
183
if persistable is not None:
184
persistable.setStyle(style)
188
def loadValueFromFile(filename, variable, passphrase=None):
189
"""Load the value of a variable in a Python file.
191
Run the contents of the file, after decrypting if C{passphrase} is
192
given, in a namespace and return the result of the variable
195
@param filename: string
196
@param variable: string
197
@param passphrase: string
203
fileObj = open(filename, mode)
204
d = {'__file__': filename}
206
data = fileObj.read()
207
data = _decrypt(passphrase, data)
214
def guessType(filename):
215
ext = os.path.splitext(filename)[1]
226
__all__ = ['loadValueFromFile', 'load', 'Persistent', 'Persistant',
227
'IPersistable', 'guessType']