228 lines
6.2 KiB
Python
228 lines
6.2 KiB
Python
|
# -*- test-case-name: twisted.test.test_sob -*-
|
||
|
|
# Copyright (c) Twisted Matrix Laboratories.
|
||
|
|
# See LICENSE for details.
|
||
|
|
|
||
|
|
#
|
||
|
|
"""
|
||
|
|
Save and load Small OBjects to and from files, using various formats.
|
||
|
|
|
||
|
|
Maintainer: Moshe Zadka
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os, sys
|
||
|
|
try:
|
||
|
|
import cPickle as pickle
|
||
|
|
except ImportError:
|
||
|
|
import pickle
|
||
|
|
try:
|
||
|
|
import cStringIO as StringIO
|
||
|
|
except ImportError:
|
||
|
|
import StringIO
|
||
|
|
from hashlib import md5
|
||
|
|
from twisted.python import log, runtime
|
||
|
|
from twisted.persisted import styles
|
||
|
|
from zope.interface import implements, Interface
|
||
|
|
|
||
|
|
# Note:
|
||
|
|
# These encrypt/decrypt functions only work for data formats
|
||
|
|
# which are immune to having spaces tucked at the end.
|
||
|
|
# All data formats which persist saves hold that condition.
|
||
|
|
def _encrypt(passphrase, data):
|
||
|
|
from Crypto.Cipher import AES as cipher
|
||
|
|
leftover = len(data) % cipher.block_size
|
||
|
|
if leftover:
|
||
|
|
data += ' '*(cipher.block_size - leftover)
|
||
|
|
return cipher.new(md5(passphrase).digest()[:16]).encrypt(data)
|
||
|
|
|
||
|
|
def _decrypt(passphrase, data):
|
||
|
|
from Crypto.Cipher import AES
|
||
|
|
return AES.new(md5(passphrase).digest()[:16]).decrypt(data)
|
||
|
|
|
||
|
|
|
||
|
|
class IPersistable(Interface):
|
||
|
|
|
||
|
|
"""An object which can be saved in several formats to a file"""
|
||
|
|
|
||
|
|
def setStyle(style):
|
||
|
|
"""Set desired format.
|
||
|
|
|
||
|
|
@type style: string (one of 'pickle' or 'source')
|
||
|
|
"""
|
||
|
|
|
||
|
|
def save(tag=None, filename=None, passphrase=None):
|
||
|
|
"""Save object to file.
|
||
|
|
|
||
|
|
@type tag: string
|
||
|
|
@type filename: string
|
||
|
|
@type passphrase: string
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
class Persistent:
|
||
|
|
|
||
|
|
implements(IPersistable)
|
||
|
|
|
||
|
|
style = "pickle"
|
||
|
|
|
||
|
|
def __init__(self, original, name):
|
||
|
|
self.original = original
|
||
|
|
self.name = name
|
||
|
|
|
||
|
|
def setStyle(self, style):
|
||
|
|
"""Set desired format.
|
||
|
|
|
||
|
|
@type style: string (one of 'pickle' or 'source')
|
||
|
|
"""
|
||
|
|
self.style = style
|
||
|
|
|
||
|
|
def _getFilename(self, filename, ext, tag):
|
||
|
|
if filename:
|
||
|
|
finalname = filename
|
||
|
|
filename = finalname + "-2"
|
||
|
|
elif tag:
|
||
|
|
filename = "%s-%s-2.%s" % (self.name, tag, ext)
|
||
|
|
finalname = "%s-%s.%s" % (self.name, tag, ext)
|
||
|
|
else:
|
||
|
|
filename = "%s-2.%s" % (self.name, ext)
|
||
|
|
finalname = "%s.%s" % (self.name, ext)
|
||
|
|
return finalname, filename
|
||
|
|
|
||
|
|
def _saveTemp(self, filename, passphrase, dumpFunc):
|
||
|
|
f = open(filename, 'wb')
|
||
|
|
if passphrase is None:
|
||
|
|
dumpFunc(self.original, f)
|
||
|
|
else:
|
||
|
|
s = StringIO.StringIO()
|
||
|
|
dumpFunc(self.original, s)
|
||
|
|
f.write(_encrypt(passphrase, s.getvalue()))
|
||
|
|
f.close()
|
||
|
|
|
||
|
|
def _getStyle(self):
|
||
|
|
if self.style == "source":
|
||
|
|
from twisted.persisted.aot import jellyToSource as dumpFunc
|
||
|
|
ext = "tas"
|
||
|
|
else:
|
||
|
|
def dumpFunc(obj, file):
|
||
|
|
pickle.dump(obj, file, 2)
|
||
|
|
ext = "tap"
|
||
|
|
return ext, dumpFunc
|
||
|
|
|
||
|
|
def save(self, tag=None, filename=None, passphrase=None):
|
||
|
|
"""Save object to file.
|
||
|
|
|
||
|
|
@type tag: string
|
||
|
|
@type filename: string
|
||
|
|
@type passphrase: string
|
||
|
|
"""
|
||
|
|
ext, dumpFunc = self._getStyle()
|
||
|
|
if passphrase:
|
||
|
|
ext = 'e' + ext
|
||
|
|
finalname, filename = self._getFilename(filename, ext, tag)
|
||
|
|
log.msg("Saving "+self.name+" application to "+finalname+"...")
|
||
|
|
self._saveTemp(filename, passphrase, dumpFunc)
|
||
|
|
if runtime.platformType == "win32" and os.path.isfile(finalname):
|
||
|
|
os.remove(finalname)
|
||
|
|
os.rename(filename, finalname)
|
||
|
|
log.msg("Saved.")
|
||
|
|
|
||
|
|
# "Persistant" has been present since 1.0.7, so retain it for compatibility
|
||
|
|
Persistant = Persistent
|
||
|
|
|
||
|
|
class _EverythingEphemeral(styles.Ephemeral):
|
||
|
|
|
||
|
|
initRun = 0
|
||
|
|
|
||
|
|
def __init__(self, mainMod):
|
||
|
|
"""
|
||
|
|
@param mainMod: The '__main__' module that this class will proxy.
|
||
|
|
"""
|
||
|
|
self.mainMod = mainMod
|
||
|
|
|
||
|
|
def __getattr__(self, key):
|
||
|
|
try:
|
||
|
|
return getattr(self.mainMod, key)
|
||
|
|
except AttributeError:
|
||
|
|
if self.initRun:
|
||
|
|
raise
|
||
|
|
else:
|
||
|
|
log.msg("Warning! Loading from __main__: %s" % key)
|
||
|
|
return styles.Ephemeral()
|
||
|
|
|
||
|
|
|
||
|
|
def load(filename, style, passphrase=None):
|
||
|
|
"""Load an object from a file.
|
||
|
|
|
||
|
|
Deserialize an object from a file. The file can be encrypted.
|
||
|
|
|
||
|
|
@param filename: string
|
||
|
|
@param style: string (one of 'pickle' or 'source')
|
||
|
|
@param passphrase: string
|
||
|
|
"""
|
||
|
|
mode = 'r'
|
||
|
|
if style=='source':
|
||
|
|
from twisted.persisted.aot import unjellyFromSource as _load
|
||
|
|
else:
|
||
|
|
_load, mode = pickle.load, 'rb'
|
||
|
|
if passphrase:
|
||
|
|
fp = StringIO.StringIO(_decrypt(passphrase,
|
||
|
|
open(filename, 'rb').read()))
|
||
|
|
else:
|
||
|
|
fp = open(filename, mode)
|
||
|
|
ee = _EverythingEphemeral(sys.modules['__main__'])
|
||
|
|
sys.modules['__main__'] = ee
|
||
|
|
ee.initRun = 1
|
||
|
|
try:
|
||
|
|
value = _load(fp)
|
||
|
|
finally:
|
||
|
|
# restore __main__ if an exception is raised.
|
||
|
|
sys.modules['__main__'] = ee.mainMod
|
||
|
|
|
||
|
|
styles.doUpgrade()
|
||
|
|
ee.initRun = 0
|
||
|
|
persistable = IPersistable(value, None)
|
||
|
|
if persistable is not None:
|
||
|
|
persistable.setStyle(style)
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
def loadValueFromFile(filename, variable, passphrase=None):
|
||
|
|
"""Load the value of a variable in a Python file.
|
||
|
|
|
||
|
|
Run the contents of the file, after decrypting if C{passphrase} is
|
||
|
|
given, in a namespace and return the result of the variable
|
||
|
|
named C{variable}.
|
||
|
|
|
||
|
|
@param filename: string
|
||
|
|
@param variable: string
|
||
|
|
@param passphrase: string
|
||
|
|
"""
|
||
|
|
if passphrase:
|
||
|
|
mode = 'rb'
|
||
|
|
else:
|
||
|
|
mode = 'r'
|
||
|
|
fileObj = open(filename, mode)
|
||
|
|
d = {'__file__': filename}
|
||
|
|
if passphrase:
|
||
|
|
data = fileObj.read()
|
||
|
|
data = _decrypt(passphrase, data)
|
||
|
|
exec data in d, d
|
||
|
|
else:
|
||
|
|
exec fileObj in d, d
|
||
|
|
value = d[variable]
|
||
|
|
return value
|
||
|
|
|
||
|
|
def guessType(filename):
|
||
|
|
ext = os.path.splitext(filename)[1]
|
||
|
|
return {
|
||
|
|
'.tac': 'python',
|
||
|
|
'.etac': 'python',
|
||
|
|
'.py': 'python',
|
||
|
|
'.tap': 'pickle',
|
||
|
|
'.etap': 'pickle',
|
||
|
|
'.tas': 'source',
|
||
|
|
'.etas': 'source',
|
||
|
|
}[ext]
|
||
|
|
|
||
|
|
__all__ = ['loadValueFromFile', 'load', 'Persistent', 'Persistant',
|
||
|
|
'IPersistable', 'guessType']
|