~mtrausch/pymcrypt/trunk

'''mcrypt wrapper for Python 2 and 3
Copyright (C) 2010 Michael B. Trausch <mike@trausch.us>
License: GPL v3 as published by the Free Software Foundation

This Python extension module makes it possible to use the mcrypt
library from within Python.  At this point it is best suited for small
encryption/decryption jobs, as it does not support (for example)
block-by-block operation just yet.  This is planned for a future
release.

PLEASE NOTE that the API provided by this module is NOT YET STABLE.
That is, the API MAY CHANGE IN A BACKWARDS INCOMPATIBLE WAY.  This is
in part because this is the first library module that I have written
for Python, and I may not get everything exactly "pythonic" the first
time around.  Furthermore, I want to make sure that the APIs used are
going to be as well-designed and non-intrusive as possible so that
picking this library up and using it in a Python program is something
that feels intuitive to any moderately experienced Python programmer.
So, please use this library for your programs, and contribute feedback
on how it can be made better.  For more information, please read the
README that is distributed with these bindings.

This wrapper attempts to expose functionality in a manner useful for
Python applications.  It does not aim to expose all of the raw
functionality of mcrypt.

Encryption keys are generated from passwords, using the PBKDF2 key
derivation function (using the OpenSSL implementation).  Optionally,
PBKDF1 can be used (for algorithms that use short enough keys) if
compatibility with other software that uses PBKDF1 is desired;
however, PBKDF1 should be considered deprecated and only used if
absolutely necessary.  Furthermore, the PBKDF1 functionality is coded
as part of this extension module.

It is of note that this wrapper does _not_ support some of the truly
advanced functionality of the mcrypt library; for example, it only
uses the compiled-in cipher algorithms and modes.  It does not permit
raw access to most of the attributes used by mcrypt.  Also, if you
provide strings it is assumed that they are UTF-8 strings; if this
assumption is not correct, or this assumption somehow breaks the way
your particular application works, then you should pass bytestrings
directly for things like salts and passwords.

Advanced functionality will be exposed by a 'raw_mcrypt' class in a
future release.'''

from cpython cimport PY_MAJOR_VERSION
from libc.stdlib cimport *
from libc.string cimport *
cimport cmcrypt as cmc

import os
import hashlib
import math
import warnings

ALGO_BLOWFISH = b'blowfish'
ALGO_DES = b'des'
ALGO_TRIPLEDES = b'tripledes'
ALGO_THREEWAY = b'threeway'
ALGO_GOST = b'gost'
ALGO_SAFERSK64 = b'safer-sk64'
ALGO_SAFERSK128 = b'safer-sk128'
ALGO_CAST128 = b'cast-128'
ALGO_XTEA = b'xtea'
ALGO_RC2 = b'rc2'
ALGO_TWOFISH = b'twofish'
ALGO_CAST256 = b'cast-256'
ALGO_SAFERPLUS = b'saferplus'
ALGO_LOKI97 = b'loki97'
ALGO_SERPENT = b'serpent'
ALGO_RIJNDAEL128 = b'rijndael-128'
ALGO_RIJNDAEL192 = b'rijndael-192'
ALGO_RIJNDAEL256 = b'rijndael-256'
ALGO_ENIGMA = b'enigma'
ALGO_ARCFOUR = b'arcfour'
ALGO_WAKE = b'wake'

ALGO_AES = ALGO_RIJNDAEL128
ALGO_DEFAULT = ALGO_AES

MODE_CBC = b'cbc'
MODE_ECB = b'ecb'
MODE_CFB = b'cfb'
MODE_OFB = b'ofb'
MODE_nOFB = b'nofb'
MODE_STREAM = b'stream'

MODE_DEFAULT = MODE_CBC

def supports_dynamic_modules():
    '''supports_dynamic_modules() -> bool
    Retrieves whether or not mcrypt supports dynamic modules.

    If the underlying mcrypt library supports dynamic modules, returns
    True; otherwise, returns False.  This function does not yet serve
    any practical use from within Python, however, because the
    functionality that mcrypt supports when dynamic modules are
    supported is not yet wrapped.'''
    return cmc.mcrypt_module_support_dynamic()

class McryptError(Exception):
    pass

class McryptWarning(UserWarning):
    pass

class McryptSecurityWarning(McryptWarning):
    pass

cdef _s2b(value):
    '''_s2b(value) -> bytes

    Returns the value as a bytestring if it is possible to do so.
    If the value is not a string, it will try to get a string from
    it before attempting conversion.  This function will always
    use UTF-8 encoding.'''
    if isinstance(value, bytes):
        return value
    elif isinstance(value, unicode):
        return value.encode('utf8')
    elif (PY_MAJOR_VERSION < 3) and isinstance(value, str):
        return value.decode('utf8')
    else:
        try:
            s = str(value)
            return _s2b(s)
        except:
            raise TypeError('value not a string and not convertable')
    return

def pbkdf1(password, salt, iterCount, keyLen, **kwargs):
    '''pbkdf1(password, salt, itercount, keyLen, **kwargs) -> bytes
    Returns an encryption key given the parameters.

    First, this is not a 100% compliant PBKDF1 function: it does not
    support MD2.  MD2 is, however, insecure, so that should not be
    much of a problem.  This function will generate up to a 16 or 20
    byte key (e.g., a key between 128 and 160 bits).  It is not the
    default KDF, but it can be specified in a call to set_parameters
    for any of the encryption or decryption objects.  It is (mostly)
    useless with newer encryption algorithms, though, which is why
    PBKDF2 is used instead by default.

    When using this function, a single keyword argument is required:
    algo, which may be one of md5 or sha1, which are the supported
    algorithms in this KDF.

    This function is not intended to be called directly from user
    code; however, it is exposed because it must be possible to
    provide it as a callback for the _mcrypt_base.set_parameters()
    method.'''
    supported = { 'md5': { 'mkl': 16, 'class': hashlib.md5 },
                  'sha1': { 'mkl': 20, 'class': hashlib.sha1 } }
    if 'algo' not in kwargs:
        raise ValueError('algo parameter must be md5 or sha1')
    if kwargs['algo'] == 'md2':
        raise ValueError('md2 is unsupported')
    if kwargs['algo'] not in supported:
        raise ValueError('algo parameter must be md5 or sha1')

    max_key_len = supported[kwargs['algo']]['mkl']
    hasher_class = supported[kwargs['algo']]['class']

    if keyLen == 0:
        keyLen = max_key_len
    elif keyLen > max_key_len:
        msg = 'max key size for algo {0} is {1}'
        raise ValueError(msg.format(kwargs['algo'], max_key_len))

    h = hasher_class()
    h.update(password + salt)
    out = h.digest()

    for i in range(iterCount - 1):
        h = hasher_class()
        h.update(out)
        out = h.digest()

    return out[:keyLen]

def pbkdf2(password, salt, iterCount, keyLen, **kwargs):
    '''pbkdf2(password, salt, iterCount, keyLen, **kwargs) -> key
    Generates an encryption key given the parameters.

    This is a wrapper aruond the OpenSSL function
    "PKCS5_PBKDF2_HMAC_SHA1", which provides the PBKDF2 functionality
    in this wrapper.  It is the default KDF.  It supports no kwargs,
    but since the KDFs have to take them, it silently ignores them.

    This function is not intended to be called directly from user
    code; however, it is exposed because it must be possible to
    provide it as a callback for the _mcrypt_base.set_parameters()
    method.'''
    cdef unsigned char *out = <unsigned char *>malloc(keyLen)
    cmc.PKCS5_PBKDF2_HMAC_SHA1(password, len(password),
                               salt, len(salt),
                               iterCount, keyLen, out)
    retval = out[:keyLen]
    free(<void *>out)

    return retval

cdef class _mcrypt_base(object):
    '''_mcrypt_base([algo[, mode]]) -> _mcrypt_base

    This is a very basic class and is not intended to be used by
    Python code.  It contains all of the functionality that can be
    used by all subclasses, as well as common data fields.  This class
    does not itself actually perform any encryption or decryption.'''

    # The underlying (native code) cipherstream object, provided by
    # the mcrypt library.  It is a pointer to an opaque C struct.
    cdef cmc.CRYPT_STREAM* _mcStream

    cdef bint _initialized

    # One of "encrypt", "decrypt", "none"
    cdef _encdec_mode

    # KDF and encryption/decryption parameters
    cdef bytes _iv
    cdef int _key_len
    cdef bytes _key
    cdef int _kdf_iterations
    cdef bytes _salt

    # object-lifetime static data
    cdef char* _algorithm
    cdef char* _mode
    cdef _avail_key_sizes
    cdef int _block_size
    cdef int _iv_len

    def __init__(self, algo = ALGO_DEFAULT, mode = MODE_DEFAULT):
        '''Initializes the base object.'''
        self._open_mcrypt(_s2b(algo), _s2b(mode))

        try:
            self._do_init()
            self._initialized = True
        except:
            self._close()

        self._algorithm = algo
        self._mode = mode

        return

    def __dealloc__(self):
        if self._mcStream is not NULL:
            cmc.mcrypt_module_close(self._mcStream)
        return

    def close(self):
        '''close() -> None

        Closes the underlying native mcrypt object.'''
        if self._initialized == False:
            raise McryptError('not open')

        self._close_mcrypt()
        return

    def get_allowed_key_sizes(self):
        '''get_allowed_key_sizes() -> list of int
        Retrieves the allowed key sizes.'''
        return self._avail_key_sizes[:]

    def get_block_size(self):
        return self._block_size

    def get_iv(self):
        '''get_iv() -> bytes
        Retrieve the current IV block, or None if none.'''
        if (self._iv_len == 0) or (len(self._iv) == 0):
            return None

        return self._iv

    def get_iv_len(self):
        '''get_iv_len() -> int
        Retrieve the IV length, or 0 if no IV needed.'''
        return self._iv_len

    def get_key(self):
        '''get_key() -> bytes
        Retrieves the encryption key.

        There is often no need to do this, so doing this is not
        recommended and will issue a McryptSecurityWarning.  Most
        likely one will only use this when attempting to debug any
        interoperability issues (if any arise) having to do with key
        generation.'''
        warnings.warn('get_key() is not recommended',
                      McryptSecurityWarning)
        klen = len(self._key)
        if klen == 0 or (klen not in self._avail_key_sizes and
                         len(self._avail_key_sizes) > 0):
            return None
        return self._key

    def get_parameters(self):
        '''get_parameters() -> dict
        Retrieves the parameters for repeat or inverse operations.

        The one parameter that is not present in the dict is the key.
        The key is assumed to be known and regenerated, and the
        regenerated key can be coupled with these parameters to, for
        example, perform a decryption when the present operation is an
        encryption.'''
        def _to_hex(bytestr):
            return ''.join(['{0:02x}'.format(x) for x in bytestr])

        retval = {}

        retval['algorithm'] = self._algorithm.decode('utf8')
        retval['mode'] = self._mode.decode('utf8')

        if self._iv is None:
            retval['iv'] = None
        elif len(self._iv) > 0:
            retval['iv'] =  _to_hex(self._iv)
        else:
            retval['iv'] = None

        retval['salt'] = _to_hex(self._salt)
        return(retval)

    def get_salt(self):
        '''Retrieve the salt value.'''
        if len(self._salt) == 0:
            return None

        return self._salt

    # XXX: function-length
    def set_parameters(self, passwd, salt = None, iv = None,
                       kdf = None, iters = 10000, klen = 0,
                       **kwargs):
        '''set_parameters(passwd[, salt[, iv[, kdf[, iters[, klen]]]]])
        Set the password for key generation.

        If salt is not specified, a new salt that is at least 8-bytes
        long will be generated.  If iv is not specified, a new IV (if
        one is required) will be generated.  If iters is not
        specified, a default of 10,000 is used.  If klen (key length)
        is not specified, the default is to use the largest key length
        permitted by the current (algorithm, mode) pairs.

        If kdf is not specified, the default key derivation function
        (PBKDF2) will be used.  Otherwise, a KDF function must be
        passed in, and the KDF must look like this:

            kdf(password, salt, iters, klen, **kwargs) -> key

        All keyword arguments are passed to the KDF, for extensibility
        purposes.

        Note that setting these parameters will rerun the KDF.  This
        can take a noticable amount of time on a slow system, or even
        on a fast system when using a very large number of iterations.
        Since we do not save the password at all, it must be provided
        every time the user wants to generate the same key.'''
        if kdf is None:
            kdf = pbkdf2
        #elif PY_MAJOR_VERSION == 2
        #    if callable(kdf) == False:
        #        raise ValueError('kdf must be callable')
        #elif PY_MAJOR_VERSION == 3:
        #    if hasattr(kdf, '__call__') == False:
        #        raise ValueError('kdf must be callable')

        if iters < 1000:
            msg = 'possibly insecure # of iterations'.format(iters)
            warnings.warn(msg, McryptSecurityWarning)
        self._kdf_iterations = iters

        if salt is not None:
            if len(_s2b(salt)) < 8:
                msg = 'extremely short salt in use'
                warnings.warn(msg, McryptSecurityWarning)
            self._salt = _s2b(salt)
        else:
            # is this secure/random enough?
            self._salt = os.urandom(8)

        if iv is not None:
            if len(_s2b(iv)) < self._iv_len:
                msg = 'iv must be {0} bytes exactly'
                raise ValueError(msg.format(self._iv_len))
            if len(_s2b(iv)) > self._iv_len:
                msg = 'IV too large, taking only first {0} bytes'
                warnings.warn(msg.format(self._iv_len))

            self._iv = _s2b(iv[:self._iv_len])
        else:
            # is this secure/random enough?
            self._iv = os.urandom(self._iv_len)

        if len(self._avail_key_sizes) == 0:
            warnings.warn('{}/{} gave no key sizes?'.
                          format(self._algorithm.decode('utf8'),
                                 self._mode.decode('utf8')))
            self._key_len = cmc.mcrypt_enc_get_key_size(self._mcStream)
            klen = self._key_len
        elif klen == 0:
            self._key_len = self._avail_key_sizes[-1]
        elif klen in self._avail_key_sizes:
            self._key_len = klen
        else:
            msg = 'klen must be one of {0}, not {1}'
            raise ValueError(msg.format(self._avail_key_sizes,
                                        klen))

        self._key = kdf(_s2b(passwd), self._salt, self._kdf_iterations,
                        self._key_len, **kwargs)

    def _close_mcrypt(self):
        if self._mcStream is not NULL:
            cmc.mcrypt_module_close(self._mcStream)
        self._mcStream = NULL
        return

    def _do_init(self):
        '''Handles variable init at object instantiation time.'''
        self._iv_len = cmc.mcrypt_enc_get_iv_size(self._mcStream)
        self._block_size = cmc.mcrypt_enc_get_block_size(self._mcStream)
        self._encdec_mode = 'none'

        self._avail_key_sizes = []
        cdef int *keyszs
        cdef int array_length

        keyszs = cmc.mcrypt_enc_get_supported_key_sizes(self._mcStream,
                                                        &array_length)
        for i in range(array_length):
            self._avail_key_sizes.append(keyszs[i])

        free(<void*>keyszs)
        return

    def _open_mcrypt(self, algo, mode):
        self._mcStream = cmc.mcrypt_module_open(algo, NULL,
                                                mode, NULL)
        if self._mcStream is NULL:
            msg = 'mcrypt init failed: a:{0} m:{1}'
            raise McryptError(msg.format(algo, mode))
        return

cdef class raw_encrypt(_mcrypt_base):
    '''raw_encrypt([algo[, mode]]) -> raw_encrypt object
    A class for encrypting data using the mcrypt library.

    This class provides encryption functionality to Python programs,
    built on the mcrypt library.  If unspecified, algo and mode will
    be set to default values, which for this release is AES
    (Rijndael-128) and CBC mode.  These defaults may change in the
    future, however.  The algorithm and the mode can be set to one of
    the constants defined in this module; see the module help for what
    those are.

    The encryption process provides a barebones encryption class; it
    does nothing but convert from plaintext to ciphertext,
    block-by-block.  Extra functionality comes from classes derived
    from raw_encrypt.

    Note that this class DOES NOT put the IV in the first block of the
    output!  This is a raw encryption class, and if you want to put
    the IV in the resulting file, you will need to do so yourself (and
    take care to remove the IV when you decrypt).  This class just
    provides the foundation for encryption using mcrypt.'''
    cdef bint _running
    cdef int _blocks_processed

    def __init__(self, algo = ALGO_DEFAULT, mode = MODE_DEFAULT):
        super(raw_encrypt, self).__init__(algo, mode)
        self._encdec_mode = 'encrypt'
        return

    def begin(self):
        '''begin() -> bool
        Begins an encryption process.

        The encryption process is a state machine: it is started,
        stepped (presumably until the end of input) and then stopped
        (whereupon it can be restarted with a different input if
        desired).  It is an error to call this method when an
        encryption process is already running.'''
        if self._running == True:
            msg = 'attempt to start an already started encrypter'
            raise McryptError(msg)

        if len(self._iv) != self._iv_len:
            raise McryptError('IV is incorrectly set.')
        elif self._iv_len == 0:
            s = cmc.mcrypt_generic_init(self._mcStream,
                                        <void*> self._key,
                                        len(self._key),
                                        NULL)
        else:
            s = cmc.mcrypt_generic_init(self._mcStream,
                                        <void*> self._key,
                                        len(self._key),
                                        <void*> self._iv)

        if s != 0:
            raise McryptError(cmc.mcrypt_strerror(s))

        self._blocks_processed = 0
        self._running = True
        return self._running

    def end(self):
        '''end() -> None
        Ends an encryption process.

        After an encryption process is ended, a new one can be begin.
        In effect, this is how object reuse is made possible.  It is
        an error to call this method when an encryption process is
        not running.'''
        if self._running == False:
            msg = 'attempt to end an already ended encrypter'
            raise McryptError(msg)

        cmc.mcrypt_generic_deinit(self._mcStream)
        return

    def get_blocks_processed(self):
        return self._blocks_processed

    def step(self, input):
        '''step(input) -> bytes
        Runs a single step of the encryption process.

        This function executes a single step of the encryption
        process; that is, it will encrypt whatever blocks of data are
        provided.  The input data length must be exactly a multiple of
        the current block size.  An exception will be raised if the
        data is incorrectly sized or if mcrypt returns an error.

        The return value is the encrypted block.'''
        if self._running == False:
            msg = 'attempt to step nonrunning encrypter'
            raise McryptError(msg)
        in_bin = _s2b(input)
        in_len = len(in_bin)

        if (in_len % self._block_size) > 0:
            print(in_len, self._block_size)
            print(in_len % self._block_size)
            msg = 'input size must be an interval of block size {0} '
            msg += '(e.g., size % {0} == 0 must be True) '
            msg += '(provided {1} % {0} = {2})'
            raise ValueError(msg.format(self._block_size, in_len,
                                        in_len % self._block_size))

        cdef char* buffer = <char*>malloc(in_len)
        memcpy(buffer, <char*>in_bin, in_len)
        success = cmc.mcrypt_generic(self._mcStream, <void*>buffer, in_len)

        if success != 0:
            free(<void*>buffer)
            msg = 'mcrypt returned error {0}: {1}'
            raise McryptError(msg.format(success,
                                         cmc.mcrypt_strerror(success)))

        retval = buffer[:in_len]
        free(<void*>buffer)

        self._blocks_processed += 1
        return retval

cdef class raw_decrypt(_mcrypt_base):
    '''raw_decrypt([algo[, mode]]) -> raw_decrypt object
    A class for decrypting data using the mcrypt library.

    This class provides decryption functionality to Python programs,
    built on the mcrypt library.  If unspecified, algo and mode will
    be set to default values, which for this release is AES
    (Rijndael-128) and CBC mode.  These defaults may change in the
    future, however.  The algorithm and the mode can be set to one of
    the constants defined in this module; see the module help for what
    those are.

    The decryption process provides a barebones decryption class.  You
    must know and set all parameters for decryption, they are not
    automatically detected.  This class just takes ciphertext blocks
    and turns them into plaintext blocks.

    Note that this class DOES NOT remove the IV from the first block
    of the output!  This is a raw decryption class, and if an IV is
    present in the file and needs to be removed, you will need to do
    so yourself.  This class just provides the foundation for
    encryption using mcrypt.'''
    cdef bint _running
    cdef int _blocks_processed

    def __init__(self, algo = ALGO_DEFAULT, mode = MODE_DEFAULT):
        super(raw_decrypt, self).__init__(algo, mode)
        self._encdec_mode = 'decrypt'
        return

    def begin(self):
        '''begin() -> bool
        Begins a decryption process.

        The decryption process is a state machine: it is started,
        stepped (presumably until the end of input) and then stopped
        (whereupon it can be restarted with a different input if
        desired).  It is an error to call this method when a
        decryption process is already running.'''
        if self._running == True:
            msg = 'attempt to start an already started decrypter'
            raise McryptError(msg)

        if len(self._iv) != self._iv_len:
            raise McryptError('IV is incorrectly set')
        elif self._iv_len == 0:
            cmc.mcrypt_generic_init(self._mcStream,
                                    <void*> self._key,
                                    len(self._key),
                                    NULL)
        else:
            cmc.mcrypt_generic_init(self._mcStream,
                                    <void*> self._key,
                                    len(self._key),
                                    <void*> self._iv)

        self._blocks_processed = 0
        self._running = True
        return self._running

    def end(self):
        '''end() -> None
        Ends a decryption process.

        After a decryption process is ended, a new one can be begin.
        In effect, this is how object reuse is made possible.  It is
        an error to call this method when a decryption process is not
        running.'''
        if self._running == False:
            msg = 'attempt to end an already ended decrypter'
            raise McryptError(msg)

        cmc.mcrypt_generic_deinit(self._mcStream)
        return

    def get_blocks_processed(self):
        return self._blocks_processed

    def step(self, input):
        '''step(input) -> bytes
        Runs a single step of the decryption process.

        This function executes a single step of the decryption
        process; that is, it will decrypt whatever blocks of data are
        provided.  The input data length must be exactly a multiple of
        the current block size.  An exception will be raised if the
        data is incorrectly sized or if mcrypt returns an error.

        The return value is the encrypted block.'''
        if self._running == False:
            msg = 'attempt to step nonrunning decrypter'
            raise McryptError(msg)
        in_bin = _s2b(input)
        in_len = len(in_bin)

        if in_len % self._block_size > 0:
            msg = 'input size must be an interval of block size {0} '
            msg += '(e.g., size % {0} == 0 must be True) '
            msg += '(provided {1} % {0} = {2})'
            raise ValueError(msg.format(self._block_size, in_len, in_len % self._block_size))

        cdef char* buffer = <char*>malloc(in_len)
        memcpy(buffer, <char*>in_bin, in_len)
        success = cmc.mdecrypt_generic(self._mcStream, <void*>buffer, in_len)

        if success != 0:
            free(<void*>buffer)
            msg = 'mcrypt returned error {0}: {1}'
            raise McryptError(msg.format(success,
                                         cmc.mcrypt_strerror(success)))

        retval = buffer[:in_len]
        free(<void*>buffer)

        self._blocks_processed += 1
        return retval