from gzip import zlib
import cStringIO
import hashlib
import hmac
import json
import os

from M2Crypto import EVP
from yoconfig.util import get_config

from authtkt.ticket import AuthTkt


class DecryptionError(Exception):
    """
    Failed to decrypt. Can happen with wrong password, for example.
    """


class EncryptionError(Exception):
    """
    Failed to encrypt. Can happen with wrong password, for example.
    """


class EncryptedAuthTkt(object):

    @staticmethod
    def from_data(secret, uid, data=None, ip='0.0.0.0', tokens=(),
                  base64=True, ts=None):
        if data is None:
            data = {}
        data = _encrypt_userdata(data)
        authticket = AuthTkt(secret, uid, data, ip, tokens, base64, ts)
        return EncryptedAuthTkt(authticket)

    def __init__(self, authticket):
        self.authticket = authticket

    @property
    def uid(self):
        return self.authticket.uid

    @property
    def data(self):
        return _decrypt_userdata(self.authticket.data)

    def ticket(self):
        return self.authticket.ticket()

    def cookie(self, name, domain=None, path='/', secure=False):
        return self.authticket.cookie(name, domain, path, secure)

    def cookie_value(self):
        return self.authticket.cookie_value()


def _encrypt_userdata(cleartext):
    cleartext = json.dumps(cleartext)
    secret = get_config('CRYPTO_SECRET')
    # the crypto algorithms are unicode unfriendly
    if isinstance(secret, unicode):
        secret = secret.encode('utf8')

    # get 256 bit random encryption salt
    salt = os.urandom(32)
    # derive 256 bit encryption key using the pbkdf2 standard
    key = EVP.pbkdf2(secret, salt, iter=1000, keylen=32)

    # Derive encryption key and HMAC key from it
    # See Practical Cryptography section 8.4.1.
    hmacKey = hashlib.sha256(key + 'MAC').digest()
    encKey = hashlib.sha256(key + 'encrypt').digest()
    del key

    # get 128 bit random iv
    iv = os.urandom(16)

    # Add HMAC to cleartext so that we can check during decrypt if we got
    # the right cleartext back. We are doing sign-then-encrypt, which let's
    # us encrypt empty cleartext (otherwise we'd need to pad with some
    # string to encrypt). Practical Cryptography by Schneier & Ferguson
    # also recommends doing it in this order in section 8.2.
    mac = hmac.new(
        hmacKey, cleartext + iv + salt, hashlib.sha256).hexdigest()
    del hmacKey

    try:
        ciphertext = _encipher(
            zlib.compress(cleartext + mac), encKey, iv, 'aes_128_ofb')
    except EVP.EVPError, e:
        raise EncryptionError(str(e))

    return (
        iv + salt + ciphertext).encode('base64').strip().replace('\n', '')


def _decrypt_userdata(ciphertext):
    ciphertext = ciphertext.decode('base64')
    iv, salt, ciphertext = (
        ciphertext[:16], ciphertext[16:48], ciphertext[48:])

    secret = get_config('CRYPTO_SECRET')
    if isinstance(secret, unicode):
        secret = secret.encode('utf8')

    # derive 256 bit key using the pbkdf2 standard
    key = EVP.pbkdf2(secret, salt, iter=1000, keylen=32)

    # Derive encryption key and HMAC key from it
    # See Practical Cryptography section 8.4.1.
    hmacKey = hashlib.sha256(key + 'MAC').digest()
    encKey = hashlib.sha256(key + 'encrypt').digest()
    del key

    # decrypt
    try:
        ret = zlib.decompress(
            _decipher(ciphertext, encKey, iv, 'aes_128_ofb'))
    except EVP.EVPError, e:
        raise DecryptionError(str(e))
    finally:
        del encKey

    # Check MAC
    mac = ret[-64:]
    ret = ret[:-64]
    try:
        if hmac.new(hmacKey, ret + iv + salt,
                    hashlib.sha256).hexdigest() != mac:
            raise DecryptionError('HMAC does not match')
    finally:
        del hmacKey

    return json.loads(ret)


def _cipherFilter(cipher, inf, outf):
    while True:
        buf = inf.read()
        if not buf:
            break
        outf.write(cipher.update(buf))
    outf.write(cipher.final())
    return outf.getvalue()


def _decipher(ciphertext, key, iv, alg):
    cipher = EVP.Cipher(alg=alg, key=key, iv=iv, op=0)
    del key
    pbuf = cStringIO.StringIO()
    cbuf = cStringIO.StringIO(ciphertext)
    plaintext = _cipherFilter(cipher, cbuf, pbuf)
    pbuf.close()
    cbuf.close()
    return plaintext


def _encipher(plaintext, key, iv, alg):
    cipher = EVP.Cipher(alg=alg, key=key, iv=iv, op=1)
    del key
    pbuf = cStringIO.StringIO(plaintext)
    cbuf = cStringIO.StringIO()
    ciphertext = _cipherFilter(cipher, pbuf, cbuf)
    pbuf.close()
    cbuf.close()
    return ciphertext

@@ -19,25 +19,8 @@
 lib%0A
-import cStringIO%0A
 impo
@@ -2465,29 +2465,16 @@
 ncipher(
-%0A            
 zlib.com
@@ -2496,16 +2496,47 @@
  + mac),
+%0A                              
  encKey,
@@ -3377,29 +3377,16 @@
 ompress(
-%0A            
 _deciphe
@@ -3822,215 +3822,96 @@
 ef _
+de
 cipher
-Filter
 (cipher
-, inf, outf):%0A    while True:%0A        buf = inf.read()%0A        if not buf:%0A            break%0A        outf.write(cipher.update(buf))%0A    outf.write(cipher.final())%0A    return outf.getvalue(
+text, key, iv, alg):%0A    return _cipherFilter(alg, key, iv, 0, ciphertext
 )%0A%0A%0A
@@ -3911,39 +3911,38 @@
 ext)%0A%0A%0Adef _
-d
 e
+n
 cipher(
-cipher
+plain
 text, key, i
@@ -3958,239 +3958,45 @@
     
-cipher = EVP.Cipher(alg=alg, key=key, iv=iv, op=0)%0A    del key%0A    pbuf = cStringIO.StringIO()%0A    cbuf = cStringIO.StringIO(ciphertext)%0A    plaintext = _cipherFilter(cipher, cbuf, pbuf)%0A    pbuf.close()%0A    cbuf.close()%0A    return
+return _cipherFilter(alg, key, iv, 1,
  pla
@@ -4001,24 +4001,25 @@
 laintext
+)
 %0A%0A%0Adef _
 encipher
@@ -4010,34 +4010,32 @@
 %0A%0A%0Adef _
-en
 cipher
-(plaintext
+Filter(alg
 , key, i
@@ -4033,27 +4033,34 @@
 g, key, iv, 
-alg
+op, input_
 ):%0A    ciphe
@@ -4102,17 +4102,18 @@
 =iv, op=
-1
+op
 )%0A    de
@@ -4126,180 +4126,50 @@
     
-pbuf = cStringIO.StringIO(plaintext)%0A    cbuf = cStringIO.StringIO()%0A    ciphertext = _cipherFilter(cipher, pbuf, cbuf)%0A    pbuf.close()%0A    cbuf.close()%0A    return ciphertext
+return cipher.update(input_) + cipher.final()
 %0A

Simplify cipher fuctions