"""
**Rijndael Symmetric Cipher**
"""
#====================================================================
# AES (Rijndael) python module
#
# Copyright (c) 2001, Bryan Mongeau <bryan@eevolved.com>
# All rights reserved.
#
# This code is hereby placed in the public domain.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1- Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2- Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# 3- The names of its contributors may not be used to endorse or promote
#    products derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#====================================================================

from rijndaelc import *

noSHA256=0
try:
  from sha256.sha256 import sha256
except ImportError:
  noSHA256=1


# Utility function for converting an arbitrary string into a hex string.
def str2hex(s):
  hex='0123456789abcdef'
  out=''
  for x in range(len(s)):
    out = out + hex[ord(s[x])>>4] + hex[ord(s[x])&0x0F]
  return out


class aes:

  """
  **Advanced Encryption Standard (aka Rijndael)**

  The aes python class wraps on the optimized ANSI C
  implementation by Paulo Barreto. 128, 192 and 256 bit keys
  are supported. The ciphering occurs with RFC2040 padding
  in ECB or CBC mode.

  **Example usage**::

    >>> from aes.aes import aes

    >>> encoder = aes()

    >>> encoder.setKey('1'*32)  # Key will be 256 bits

    >>> encoder.encrypt("uncrackable")

    '\300\017:\243\204r\330\263\350O>\361\007\2439\370'

    >>> code = encoder.encrypt(''uncrackable'')

    >>> encoder.decrypt(code)

    ''uncrackable''

    >>> encoder.setKey('2'*64)  # switch to another key

    >>> encoder.decrypt(code)

    *Traceback (most recent call last):*

      File "rijndael.py", line 138, in decrypt

      raise RuntimeError("Could not decrypt")

    RuntimeError: Could not decrypt

  """

  def __init__(self, key=None, mode="ECB", IV='' ):

    """
    Cipher Initialisation

    If a *key* string is provided, it will be used immediately
    to setup internal encryption and decryption keys. Keys can be either
    128, 192 or 256 bits, thus the key string ought to either 16, 24 or 32
    bytes in length.

    The cipher *mode* ECB is set by default. CBC mode can also be used
    in conjunction with the initialization vector *IV*. The initialization
    vector can be from 0 to 16 bytes in length.

    The strength of any cipher depends on the probability
    of guessing its key. Please note that for real world
    crypto applications in a hostile environment, make sure
    that your "random" keys are generated by a cryptographically
    secure pseudo-random number generator (CSPRNG) or better.
    """

    self.encKey = new_keyInstance()
    self.decKey = new_keyInstance()
    self.cipher = new_cipherInstance()

    if not key:
      self.keyMaterial = None
    else:
      self.setKey(key)

    self.setCipherMode(mode,IV)


  def setCipherMode(self, mode, IV=''):
    """
    Sets the cipher mode to use.

    *mode* can be either "ECB" or "CBC". If CBC, you can provide
    the iniatialization vector *IV*.

    ECB stands for Electronic Code Book and encrypts blocks of data
    serially. Simple and fast, but repeating patterns will produce
    repeating encrypted data, not a Good Thing in the face of cryptanalysis.
    CBC stands for cipher block chaining mode and it chops up encryption
    blocks in an attempt to mask repeating patterns.
    """

    if mode=="CBC":
      if len(IV) > 16:
        raise RuntimeError("Initialization Vector exceeds 16 bytes.")
      else:
        cipherInit(self.cipher, 2, IV)

    elif mode=="ECB":
      cipherInit(self.cipher, 1, '')

    else:
      raise RuntimeError('Cipher mode must be "ECB" or "CBC"')


  def setKey(self, key):

    """
    Set the key you wish to use for encryption / decryption.

    Pretty simple. The key must adhere to the format defined in the
    constructor of the class.
    """

    if len(key)==16 or len(key)==24 or len(key)==32:

      makeKey(self.encKey,0,len(key)*8, str2hex(key))
      makeKey(self.decKey,1,len(key)*8, str2hex(key))
      self.keyMaterial = key

    else:
      raise RuntimeError,"Invalid Key, not 16 or 24 or 32 bytes."


  def encrypt(self, data):

    """
      - Will accept any python string (NULL bytes or no)
        and will RFC2040 pad it (16 byte pad) to be a nice block size.

      - No length checking is performed on the inbound data, so beware
        in extreme cases (I tested it up to 100 MB). For very large data
        lengths, consider processing in batches to conserve memory.

    The encrypted string (with NULL bytes) is returned.
    """

    if str(type(data)) != "<type 'string'>":
      raise TypeError('Data to encrypt must be a string')

    if not self.keyMaterial:
      raise RuntimeError("No encryption key")

    try:
      return padEncrypt(self.cipher, self.encKey, data, len(data) )
    except:
      raise RuntimeError("Could not encrypt")


  def decrypt(self, data):

    """
      - Encrypted data is ECB or CBC decoded in 16 byte blocks and unpadded.

      - The unencrypted result is returned. If it looks like garbage to
        you, check the key you're using!
    """

    if str(type(data)) != "<type 'string'>":
      raise TypeError('Data to decrypt must be a string')

    if not self.keyMaterial:
      raise RuntimeError("No decryption key")

    try:
      return padDecrypt(self.cipher, self.decKey, data, len(data) )
    except:
      raise RuntimeError("Could not decrypt")


  def lazyEncrypt(self, passphrase, data):

    """
    What a lazy encrypt does is generate a 256 bit key to use for ciphering
    based on a passphrase of arbitrary length.  The passphrase is hashed with
    SHA256 and the data is encrypted with the digest.

    Makes password-based encryption simple.

    Note: The sha256 module is required.
    """
    if noSHA256:
      raise RuntimeError("Need sha256 module to lazyEncrypt")

    if str(type(data)) != "<type 'string'>":
      raise TypeError('Data to encrypt must be a string')

    hasher = sha256()
    hasher.update(passphrase)
    key = str2hex(hasher.digest())
    tmpKey = new_keyInstance()
    makeKey( tmpKey, 0, 256, key)

    try:
      return padEncrypt(self.cipher, tmpKey, data, len(data))
    except:
      raise RuntimeError("Could not lazyEncrypt")


  def lazyDecrypt(self, passphrase, data):
    """
    Decrypt a block of data that was encrypted with a passphrase.

    The passphrase can be of arbitrary length: it is SHA256 hashed
    to a uniform size (256 bits). The resulting digest is used as the
    deciphering key.

    If the result is gibberish, check the passphrase validity.

    Note: The sha256 module is required.
    """
    if noSHA256:
      raise RuntimeError("Need sha256 module to lazyDecrypt")

    if str(type(data)) != "<type 'string'>":
      raise TypeError('Data to decrypt must be a string')

    hasher = sha256()
    hasher.update(passphrase)
    key = str2hex(hasher.digest())
    tmpKey = new_keyInstance()
    makeKey( tmpKey, 1, 256, key)
    try:
      return padDecrypt(self.cipher, tmpKey, data, len(data))
    except:
      raise RuntimeError("Could not lazyDecrypt")
