2
0
Эх сурвалжийг харах

crypto: convert attic.crypto into a cython module

This will give a slight performance boost and make the openssl
detection/linking more robust.
Jonas Borgström 11 жил өмнө
parent
commit
0e39acffd3
4 өөрчлөгдсөн 139 нэмэгдсэн , 95 устгасан
  1. 0 89
      attic/crypto.py
  2. 111 0
      attic/crypto.pyx
  3. 2 2
      attic/testsuite/crypto.py
  4. 26 4
      setup.py

+ 0 - 89
attic/crypto.py

@@ -1,89 +0,0 @@
-"""A thin ctypes based wrapper for OpenSSL 1.0
-"""
-import os
-from ctypes import cdll, c_char_p, c_int, c_uint, c_void_p, POINTER, create_string_buffer
-from ctypes.util import find_library
-import struct
-
-
-def _find_libcrypto():
-    _possible_paths = [
-        find_library('crypto'),
-        os.environ.get('ATTIC_LIBCRYPTO_PATH'),
-        '/usr/local/opt/openssl/lib/libcrypto.dylib',  # OS X Brew
-        '/usr/local/lib/libcrypto.so',                 # FreeBSD Ports
-        '/usr/local/ssl/lib/libcrypto.so'
-    ]
-    for path in _possible_paths:
-        try:
-            lib = cdll.LoadLibrary(path)
-            if hasattr(lib, 'PKCS5_PBKDF2_HMAC'):
-                return lib
-        except OSError:
-            pass
-    raise Exception('Failed to find libcrypto version >= 1.0')
-
-libcrypto = _find_libcrypto()
-
-libcrypto.PKCS5_PBKDF2_HMAC.argtypes = (c_char_p, c_int, c_char_p, c_int, c_int, c_void_p, c_int, c_char_p)
-libcrypto.EVP_sha256.restype = c_void_p
-libcrypto.AES_set_encrypt_key.argtypes = (c_char_p, c_int, c_char_p)
-libcrypto.AES_ctr128_encrypt.argtypes = (c_char_p, c_char_p, c_int, c_char_p, c_char_p, c_char_p, POINTER(c_uint))
-libcrypto.RAND_bytes.argtypes = (c_char_p, c_int)
-libcrypto.RAND_bytes.restype = c_int
-
-_int = struct.Struct('>I')
-_long = struct.Struct('>Q')
-
-bytes_to_int = lambda x, offset=0: _int.unpack_from(x, offset)[0]
-bytes_to_long = lambda x, offset=0: _long.unpack_from(x, offset)[0]
-long_to_bytes = lambda x: _long.pack(x)
-
-
-def num_aes_blocks(length):
-    """Return the number of AES blocks required to encrypt/decrypt *length* bytes of data
-    """
-    return (length + 15) // 16
-
-
-def pbkdf2_sha256(password, salt, iterations, size):
-    """Password based key derivation function 2 (RFC2898)
-    """
-    key = create_string_buffer(size)
-    rv = libcrypto.PKCS5_PBKDF2_HMAC(password, len(password), salt, len(salt), iterations, libcrypto.EVP_sha256(), size, key)
-    if not rv:
-        raise Exception('PKCS5_PBKDF2_HMAC failed')
-    return key.raw
-
-
-def get_random_bytes(n):
-    """Return n cryptographically strong pseudo-random bytes
-    """
-    buf = create_string_buffer(n)
-    if libcrypto.RAND_bytes(buf, n) < 1:
-        raise Exception('RAND_bytes failed')
-    return buf.raw
-
-
-class AES:
-    """A thin wrapper around the OpenSSL AES CTR_MODE cipher
-    """
-    def __init__(self, key, iv=None):
-        self.key = create_string_buffer(2000)
-        self.iv = create_string_buffer(16)
-        self.buf = create_string_buffer(16)
-        self.num = c_uint()
-        self.reset(key, iv)
-
-    def reset(self, key=None, iv=None):
-        if key:
-            libcrypto.AES_set_encrypt_key(key, len(key) * 8, self.key)
-        if iv:
-            self.iv.raw = iv
-        self.num.value = 0
-
-    def encrypt(self, data):
-        out = create_string_buffer(len(data))
-        libcrypto.AES_ctr128_encrypt(data, out, len(data), self.key, self.iv, self.buf, self.num)
-        return out.raw
-    decrypt = encrypt

+ 111 - 0
attic/crypto.pyx

@@ -0,0 +1,111 @@
+"""A thin OpenSSL wrapper
+
+This could be replaced by PyCrypto or something similar when the performance
+of their PBKDF2 implementation is comparable to the OpenSSL version.
+"""
+
+from libc.string cimport memcpy
+from libc.stdlib cimport malloc, free
+
+cdef extern from "openssl/rand.h":
+    int  RAND_bytes(unsigned char *buf,int num)
+
+cdef extern from "openssl/aes.h":
+    ctypedef struct AES_KEY:
+        pass
+
+    int AES_set_encrypt_key(const unsigned char *userKey, const int bits, AES_KEY *key)
+    void AES_ctr128_encrypt(const unsigned char *in_, unsigned char *out,
+                            size_t length, const AES_KEY *key,
+                            unsigned char *ivec,
+                            unsigned char *ecount_buf,
+                            unsigned int *num)
+
+cdef extern from "openssl/evp.h":
+    ctypedef struct EVP_MD:
+        pass
+    const EVP_MD *EVP_sha256()
+    int PKCS5_PBKDF2_HMAC(const char *password, int passwordlen,
+                          const unsigned char *salt, int saltlen, int iter,
+                          const EVP_MD *digest,
+                          int keylen, unsigned char *out)
+
+import struct
+
+_int = struct.Struct('>I')
+_long = struct.Struct('>Q')
+
+bytes_to_int = lambda x, offset=0: _int.unpack_from(x, offset)[0]
+bytes_to_long = lambda x, offset=0: _long.unpack_from(x, offset)[0]
+long_to_bytes = lambda x: _long.pack(x)
+
+
+def num_aes_blocks(length):
+    """Return the number of AES blocks required to encrypt/decrypt *length* bytes of data
+    """
+    return (length + 15) // 16
+
+
+def pbkdf2_sha256(password, salt, iterations, size):
+    """Password based key derivation function 2 (RFC2898)
+    """
+    cdef unsigned char *key = <unsigned char *>malloc(size)
+    if not key:
+        raise MemoryError
+    try:
+        rv = PKCS5_PBKDF2_HMAC(password, len(password), salt, len(salt), iterations, EVP_sha256(), size, key)
+        if not rv:
+            raise Exception('PKCS5_PBKDF2_HMAC failed')
+        return key[:size]
+    finally:
+        free(key)
+
+
+def get_random_bytes(n):
+    """Return n cryptographically strong pseudo-random bytes
+    """
+    cdef unsigned char *buf = <unsigned char *>malloc(n)
+    if not buf:
+        raise MemoryError
+    try:
+        if RAND_bytes(buf, n) < 1:
+            raise Exception('RAND_bytes failed')
+        return buf[:n]
+    finally:
+        free(buf)
+
+
+cdef class AES:
+    """A thin wrapper around the OpenSSL AES CTR_MODE cipher
+    """
+    cdef AES_KEY key
+    cdef unsigned char _iv[16]
+    cdef unsigned char buf[16]
+    cdef unsigned int num
+
+    def __cinit__(self, key, iv=None):
+        self.reset(key, iv)
+
+    def reset(self, key=None, iv=None):
+        if key:
+            AES_set_encrypt_key(key, len(key) * 8, &self.key)
+        if iv:
+            memcpy(self._iv, <unsigned char *>iv, 16)
+        self.num = 0
+
+    @property
+    def iv(self):
+        return self._iv[:16]
+
+    def encrypt(self, data):
+        cdef int n = len(data)
+        cdef unsigned char *out = <unsigned char *>malloc(n)
+        if not out:
+            raise MemoryError
+        try:
+            AES_ctr128_encrypt(data, out, len(data), &self.key, self._iv, self.buf, &self.num)
+            return out[:n]
+        finally:
+            free(out)
+    decrypt = encrypt
+

+ 2 - 2
attic/testsuite/crypto.py

@@ -31,10 +31,10 @@ class CryptoTestCase(AtticTestCase):
         key = b'X' * 32
         data = b'foo' * 10
         aes = AES(key)
-        self.assert_equal(bytes_to_long(aes.iv.raw, 8), 0)
+        self.assert_equal(bytes_to_long(aes.iv, 8), 0)
         cdata = aes.encrypt(data)
         self.assert_equal(hexlify(cdata), b'c6efb702de12498f34a2c2bbc8149e759996d08bf6dc5c610aefc0c3a466')
-        self.assert_equal(bytes_to_long(aes.iv.raw, 8), 2)
+        self.assert_equal(bytes_to_long(aes.iv, 8), 2)
         self.assert_not_equal(data, aes.decrypt(cdata))
         aes.reset(iv=b'\0' * 16)
         self.assert_equal(data, aes.decrypt(cdata))

+ 26 - 4
setup.py

@@ -2,7 +2,6 @@
 import os
 import sys
 from glob import glob
-import attic
 
 import versioneer
 versioneer.versionfile_source = 'attic/_version.py'
@@ -21,6 +20,7 @@ try:
 except ImportError:
     from distutils.core import setup, Extension
 
+crypto_source = 'attic/crypto.pyx'
 chunker_source = 'attic/chunker.pyx'
 hashindex_source = 'attic/hashindex.pyx'
 
@@ -36,7 +36,7 @@ try:
             versioneer.cmd_sdist.__init__(self, *args, **kwargs)
 
         def make_distribution(self):
-            self.filelist.extend(['attic/chunker.c', 'attic/_chunker.c', 'attic/hashindex.c', 'attic/_hashindex.c'])
+            self.filelist.extend(['attic/crypto.c', 'attic/chunker.c', 'attic/_chunker.c', 'attic/hashindex.c', 'attic/_hashindex.c'])
             super(Sdist, self).make_distribution()
 
 except ImportError:
@@ -44,11 +44,32 @@ except ImportError:
         def __init__(self, *args, **kwargs):
             raise Exception('Cython is required to run sdist')
 
+    crypto_source = crypto_source.replace('.pyx', '.c')
     chunker_source = chunker_source.replace('.pyx', '.c')
     hashindex_source = hashindex_source.replace('.pyx', '.c')
     from distutils.command.build_ext import build_ext
-    if not os.path.exists(chunker_source) or not os.path.exists(hashindex_source):
-        raise ImportError('The GIT version of attic needs Cython. Install Cython or use a released version')
+    if not all(os.path.exists(path) for path in [crypto_source, chunker_source, hashindex_source]):
+        raise ImportError('The GIT version of Attic needs Cython. Install Cython or use a released version')
+
+
+def detect_openssl(prefixes):
+    for prefix in prefixes:
+        filename = os.path.join(prefix, 'include', 'openssl', 'evp.h')
+        if os.path.exists(filename):
+            with open(filename, 'r') as fd:
+                if 'PKCS5_PBKDF2_HMAC(' in fd.read():
+                    return prefix
+
+
+possible_openssl_prefixes = ['/usr', '/usr/local', '/usr/local/opt/openssl', '/usr/local/ssl', '/usr/local/openssl', '/usr/local/attic']
+if os.environ.get('ATTIC_OPENSSL_PREFIX'):
+    possible_openssl_prefixes.insert(0, os.environ.get('ATTIC_OPENSSL_PREFIX'))
+ssl_prefix = detect_openssl(possible_openssl_prefixes)
+if not ssl_prefix:
+    raise Exception('Unable to find OpenSSL >= 1.0 headers. (Looked here: {})'.format(', '.join(possible_openssl_prefixes)))
+include_dirs = [os.path.join(ssl_prefix, 'include')]
+library_dirs = [os.path.join(ssl_prefix, 'lib')]
+
 
 with open('README.rst', 'r') as fd:
     long_description = fd.read()
@@ -82,6 +103,7 @@ setup(
     scripts=['scripts/attic'],
     cmdclass=cmdclass,
     ext_modules=[
+        Extension('attic.crypto', [crypto_source], libraries=['crypto'], include_dirs=include_dirs, library_dirs=library_dirs),
         Extension('attic.chunker', [chunker_source]),
         Extension('attic.hashindex', [hashindex_source])
     ],