# -*- 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']