Initial AES-NI support

9 files changed, 462 insertions, 5 deletions

diff --git a/configure b/configure
index 89592b5..02dfadc 100755
--- a/configure
+++ b/configure
@@ -4373,6 +4373,39 @@ else
 fi
 
 
+# Checks for compiler flags.
+old_CFLAGS="$CFLAGS"
+CFLAGS=-maes
+{ $as_echo "$as_me:${as_lineno-$LINENO}: checking whether CC supports -maes" >&5
+$as_echo_n "checking whether CC supports -maes... " >&6; }
+cat confdefs.h - <<_ACEOF >conftest.$ac_ext
+/* end confdefs.h.  */
+
+int
+main ()
+{
+
+  ;
+  return 0;
+}
+_ACEOF
+if ac_fn_c_try_compile "$LINENO"; then :
+
+  { $as_echo "$as_me:${as_lineno-$LINENO}: result: yes" >&5
+$as_echo "yes" >&6; }
+
+$as_echo "#define HAVE_MAES 1" >>confdefs.h
+
+  have_maes=yes
+
+else
+  { $as_echo "$as_me:${as_lineno-$LINENO}: result: no" >&5
+$as_echo "no" >&6; }
+
+fi
+rm -f core conftest.err conftest.$ac_objext conftest.$ac_ext
+CFLAGS="$old_cflags"
+
 # Checks for libraries.
 
 # Check whether --with-gmp was given.
@@ -4921,7 +4954,7 @@ fi
 done
 
 
-for ac_header in inttypes.h sys/inttypes.h limits.h stddef.h stdint.h stdlib.h string.h wchar.h
+for ac_header in inttypes.h sys/inttypes.h cpuid.h limits.h stddef.h stdint.h stdlib.h string.h wchar.h
 do :
   as_ac_Header=`$as_echo "ac_cv_header_$ac_header" | $as_tr_sh`
 ac_fn_c_check_header_mongrel "$LINENO" "$ac_header" "$as_ac_Header" "$ac_includes_default"
@@ -4934,6 +4967,25 @@ fi
 
 done
 
+old_CPPFLAGS="$CPPFLAGS"
+if test "x$have_maes" = "xyes"; then :
+
+    CPPFLAGS="$CPPFLAGS -maes"
+
+fi
+for ac_header in wmmintrin.h
+do :
+  ac_fn_c_check_header_mongrel "$LINENO" "wmmintrin.h" "ac_cv_header_wmmintrin_h" "$ac_includes_default"
+if test "x$ac_cv_header_wmmintrin_h" = xyes; then :
+  cat >>confdefs.h <<_ACEOF
+#define HAVE_WMMINTRIN_H 1
+_ACEOF
+
+fi
+
+done
+
+CPPFLAGS="$old_CPPFLAGS"
 
 # Checks for typedefs, structures, and compiler characteristics.
 { $as_echo "$as_me:${as_lineno-$LINENO}: checking for inline" >&5
diff --git a/configure.ac b/configure.ac
index 5541361..16078ac 100644
--- a/configure.ac
+++ b/configure.ac
@@ -49,6 +49,18 @@ AX_CHECK_COMPILE_FLAG([-Wextra], [AX_APPEND_FLAG([-Wextra])])
 AX_CHECK_COMPILE_FLAG([-Wno-missing-field-initializers], [AX_APPEND_FLAG([-Wno-missing-field-initializers])])
 AX_CHECK_COMPILE_FLAG([-Wno-unused-parameter], [AX_APPEND_FLAG([-Wno-unused-parameter])])
 
+# Checks for compiler flags.
+old_CFLAGS="$CFLAGS"
+CFLAGS=-maes
+AC_MSG_CHECKING([whether CC supports -maes])
+AC_COMPILE_IFELSE([AC_LANG_PROGRAM([])], [
+  AC_MSG_RESULT([yes])
+  AC_DEFINE(HAVE_MAES, [1], [Define if CC supports -maes])
+  have_maes=yes
+  ], [AC_MSG_RESULT([no])]
+)
+CFLAGS="$old_cflags"
+
 # Checks for libraries.
 AC_ARG_WITH([gmp], AS_HELP_STRING([--without-gmp], [Build without gmp library (default: test)]))
 AS_IF([test "x$with_gmp" != "xno"], [
@@ -76,7 +88,13 @@ AC_CHECK_DECLS([mpz_powm_sec], [], [], [
 ]])
 
 # Checks for header files.
-AC_CHECK_HEADERS([inttypes.h sys/inttypes.h limits.h stddef.h stdint.h stdlib.h string.h wchar.h])
+AC_CHECK_HEADERS([inttypes.h sys/inttypes.h cpuid.h limits.h stddef.h stdint.h stdlib.h string.h wchar.h])
+old_CPPFLAGS="$CPPFLAGS"
+AS_IF([test "x$have_maes" = "xyes"], [
+    CPPFLAGS="$CPPFLAGS -maes"
+])
+AC_CHECK_HEADERS([wmmintrin.h])
+CPPFLAGS="$old_CPPFLAGS"
 
 # Checks for typedefs, structures, and compiler characteristics.
 AC_C_INLINE
diff --git a/lib/Crypto/Cipher/AES.py b/lib/Crypto/Cipher/AES.py
index eb91c5f..2aaaaa8 100644
--- a/lib/Crypto/Cipher/AES.py
+++ b/lib/Crypto/Cipher/AES.py
@@ -48,6 +48,16 @@ __revision__ = "$Id$"
 
 from Crypto.Cipher import blockalgo
 from Crypto.Cipher import _AES
+from Crypto.Util import cpuid
+# Import _AESNI. If AES-NI is not available or _AESNI has not been built, set
+# _AESNI to None.
+try:
+    if cpuid.have_aes_ni():
+        from Crypto.Cipher import _AESNI
+    else:
+        _AESNI = None
+except ImportError:
+    _AESNI = None
 
 class AESCipher (blockalgo.BlockAlgo):
     """AES cipher object"""
@@ -56,7 +66,18 @@ class AESCipher (blockalgo.BlockAlgo):
         """Initialize an AES cipher object
         
         See also `new()` at the module level."""
-        blockalgo.BlockAlgo.__init__(self, _AES, key, *args, **kwargs)
+
+        # Check if the use_aesni was specified.
+        use_aesni = True
+        if 'use_aesni' in kwargs:
+            use_aesni = kwargs['use_aesni']
+            del kwargs['use_aesni']
+
+        # Use _AESNI if the user requested AES-NI and it's available
+        if _AESNI is not None and use_aesni:
+            blockalgo.BlockAlgo.__init__(self, _AESNI, key, *args, **kwargs)
+        else:
+            blockalgo.BlockAlgo.__init__(self, _AES, key, *args, **kwargs)
 
 def new(key, *args, **kwargs):
     """Create a new AES cipher
@@ -88,6 +109,8 @@ def new(key, *args, **kwargs):
         (*Only* `MODE_CFB`).The number of bits the plaintext and ciphertext
         are segmented in.
         It must be a multiple of 8. If 0 or not specified, it will be assumed to be 8.
+      use_aesni : boolean
+        Use AES-NI if available.
 
     :Return: an `AESCipher` object
     """
diff --git a/lib/Crypto/SelfTest/Cipher/common.py b/lib/Crypto/SelfTest/Cipher/common.py
index 8bebed9..26fabed 100644
--- a/lib/Crypto/SelfTest/Cipher/common.py
+++ b/lib/Crypto/SelfTest/Cipher/common.py
@@ -289,7 +289,7 @@ class IVLengthTest(unittest.TestCase):
     def _dummy_counter(self):
         return "\0" * self.module.block_size
 
-def make_block_tests(module, module_name, test_data):
+def make_block_tests(module, module_name, test_data, additional_params=dict()):
     tests = []
     extra_tests_added = 0
     for i in range(len(test_data)):
@@ -326,6 +326,7 @@ def make_block_tests(module, module_name, test_data):
         name = "%s #%d: %s" % (module_name, i+1, description)
         params['description'] = name
         params['module_name'] = module_name
+        params.update(additional_params)
 
         # Add extra test(s) to the test suite before the current test
         if not extra_tests_added:
diff --git a/lib/Crypto/SelfTest/Cipher/test_AES.py b/lib/Crypto/SelfTest/Cipher/test_AES.py
index ea7c323..a675f8e 100644
--- a/lib/Crypto/SelfTest/Cipher/test_AES.py
+++ b/lib/Crypto/SelfTest/Cipher/test_AES.py
@@ -1422,8 +1422,14 @@ test_data = [
 
 def get_tests(config={}):
     from Crypto.Cipher import AES
+    from Crypto.Util import cpuid
     from common import make_block_tests
-    return make_block_tests(AES, "AES", test_data)
+
+    tests = make_block_tests(AES, "AES", test_data, {'use_aesni': False})
+    if cpuid.have_aes_ni():
+        # Run tests with AES-NI instructions if they are available.
+        tests += make_block_tests(AES, "AESNI", test_data, {'use_aesni': True})
+    return tests
 
 if __name__ == '__main__':
     import unittest
diff --git a/setup.py b/setup.py
index 85bd65a..5af83ac 100644
--- a/setup.py
+++ b/setup.py
@@ -162,6 +162,21 @@ class PCTBuildExt (build_ext):
                 "Crypto.PublicKey._fastmath.")
             self.__remove_extensions(["Crypto.PublicKey._fastmath"])
 
+        # Detect if we have AES-NI instrincs available
+        if not ac.get("HAVE_WMMINTRIN_H"):
+            # AES-NI instrincs not available
+            self.__remove_extensions(["Crypto.Cipher._AESNI"])
+        elif ac.get("HAVE_MAES"):
+            # -maes has to be passed to the compiler to use the AES-NI instrincs
+            self.__add_extension_compile_option(["Crypto.Cipher._AESNI"],
+                                                ["-maes"])
+
+    def __add_extension_compile_option(self, names, options):
+        """Add compiler options for the specified extension(s)"""
+        for extension in self.extensions:
+            if extension.name in names:
+                extension.extra_compile_args = options
+
     def __add_extension_link_option(self, names, options):
         """Add linker options for the specified extension(s)"""
         i = 0
@@ -426,6 +441,9 @@ kw = {'name':"pycrypto",
             Extension("Crypto.Cipher._AES",
                       include_dirs=['src/'],
                       sources=["src/AES.c"]),
+            Extension("Crypto.Cipher._AESNI",
+                      include_dirs=['src/'],
+                      sources=["src/AESNI.c"]),
             Extension("Crypto.Cipher._ARC2",
                       include_dirs=['src/'],
                       sources=["src/ARC2.c"]),
@@ -454,6 +472,9 @@ kw = {'name':"pycrypto",
             Extension("Crypto.Util.strxor",
                       include_dirs=['src/'],
                       sources=['src/strxor.c']),
+            Extension("Crypto.Util.cpuid",
+                      include_dirs=['src/'],
+                      sources=['src/cpuid.c']),
 
             # Counter modules
             Extension("Crypto.Util._counter",
diff --git a/src/AESNI.c b/src/AESNI.c
new file mode 100644
index 0000000..bd80581
--- /dev/null
+++ b/src/AESNI.c
@@ -0,0 +1,227 @@
+/*
+ *  AESNI.c: AES using AES-NI instructions
+ *
+ * Written in 2013 by Sebastian Ramacher <sebastian@ramacher.at>
+ *
+ * ===================================================================
+ * The contents of this file are dedicated to the public domain.  To
+ * the extent that dedication to the public domain is not available,
+ * everyone is granted a worldwide, perpetual, royalty-free,
+ * non-exclusive license to exercise all rights associated with the
+ * contents of this file for any purpose whatsoever.
+ * No rights are reserved.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
+ * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
+ * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+ * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * ===================================================================
+ */
+
+#include <assert.h>
+#include <stdlib.h>
+#include "Python.h"
+#include <wmmintrin.h>
+
+#define MODULE_NAME _AESNI
+#define BLOCK_SIZE 16
+#define KEY_SIZE 0
+
+#define MAXKC	(256/32)
+#define MAXKB	(256/8)
+#define MAXNR	14
+
+typedef unsigned char u8;
+
+typedef struct {
+	__m128i ek[MAXNR + 1];
+	__m128i dk[MAXNR + 1];
+	int rounds;
+} block_state;
+
+/* Helper functions to expand keys */
+
+static __m128i aes128_keyexpand(__m128i key, __m128i keygened, int shuf)
+{
+    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
+    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
+    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
+    keygened = _mm_shuffle_epi32(keygened, shuf);
+    return _mm_xor_si128(key, keygened);
+}
+
+static __m128i aes192_keyexpand_2(__m128i key, __m128i key2)
+{
+    key = _mm_shuffle_epi32(key, 0xff);
+    key2 = _mm_xor_si128(key2, _mm_slli_si128(key2, 4));
+    return _mm_xor_si128(key, key2);
+}
+
+#define KEYEXP128(K, I) aes128_keyexpand(K, _mm_aeskeygenassist_si128(K, I), 0xff)
+#define KEYEXP192(K1, K2, I) aes128_keyexpand(K1, _mm_aeskeygenassist_si128(K2, I), 0x55)
+#define KEYEXP192_2(K1, K2) aes192_keyexpand_2(K1, K2)
+#define KEYEXP256(K1, K2, I) aes128_keyexpand(K1, _mm_aeskeygenassist_si128(K2, I), 0xff)
+#define KEYEXP256_2(K1, K2) aes128_keyexpand(K1, _mm_aeskeygenassist_si128(K2, 0x00), 0xaa)
+
+/* Encryption key setup */
+static void aes_key_setup_enc(__m128i rk[], const u8* cipherKey, int keylen)
+{
+    switch (keylen) {
+        case 16:
+        {
+            /* 128 bit key setup */
+            rk[0] = _mm_loadu_si128((const __m128i*) cipherKey);
+            rk[1] = KEYEXP128(rk[0], 0x01);
+            rk[2] = KEYEXP128(rk[1], 0x02);
+            rk[3] = KEYEXP128(rk[2], 0x04);
+            rk[4] = KEYEXP128(rk[3], 0x08);
+            rk[5] = KEYEXP128(rk[4], 0x10);
+            rk[6] = KEYEXP128(rk[5], 0x20);
+            rk[7] = KEYEXP128(rk[6], 0x40);
+            rk[8] = KEYEXP128(rk[7], 0x80);
+            rk[9] = KEYEXP128(rk[8], 0x1B);
+            rk[10] = KEYEXP128(rk[9], 0x36);
+            break;
+        }
+        case 24:
+        {
+            /* 192 bit key setup */
+            __m128i temp[2];
+            rk[0] = _mm_loadu_si128((const __m128i*) cipherKey);
+            rk[1] = _mm_loadu_si128((const __m128i*) (cipherKey+16));
+            temp[0] = KEYEXP192(rk[0], rk[1], 0x01);
+            temp[1] = KEYEXP192_2(temp[0], rk[1]);
+            rk[1] = (__m128i)_mm_shuffle_pd((__m128d)rk[1], (__m128d)temp[0], 0);
+            rk[2] = (__m128i)_mm_shuffle_pd((__m128d)temp[0], (__m128d)temp[1], 1);
+            rk[3] = KEYEXP192(temp[0], temp[1], 0x02);
+            rk[4] = KEYEXP192_2(rk[3], temp[1]);
+            temp[0] = KEYEXP192(rk[3], rk[4], 0x04);
+            temp[1] = KEYEXP192_2(temp[0], rk[4]);
+            rk[4] = (__m128i)_mm_shuffle_pd((__m128d)rk[4], (__m128d)temp[0], 0);
+            rk[5] = (__m128i)_mm_shuffle_pd((__m128d)temp[0], (__m128d)temp[1], 1);
+            rk[6] = KEYEXP192(temp[0], temp[1], 0x08);
+            rk[7] = KEYEXP192_2(rk[6], temp[1]);
+            temp[0] = KEYEXP192(rk[6], rk[7], 0x10);
+            temp[1] = KEYEXP192_2(temp[0], rk[7]);
+            rk[7] = (__m128i)_mm_shuffle_pd((__m128d)rk[7], (__m128d)temp[0], 0);
+            rk[8] = (__m128i)_mm_shuffle_pd((__m128d)temp[0], (__m128d)temp[1], 1);
+            rk[9] = KEYEXP192(temp[0], temp[1], 0x20);
+            rk[10] = KEYEXP192_2(rk[9], temp[1]);
+            temp[0] = KEYEXP192(rk[9], rk[10], 0x40);
+            temp[1] = KEYEXP192_2(temp[0], rk[10]);
+            rk[10] = (__m128i)_mm_shuffle_pd((__m128d)rk[10], (__m128d) temp[0], 0);
+            rk[11] = (__m128i)_mm_shuffle_pd((__m128d)temp[0],(__m128d) temp[1], 1);
+            rk[12] = KEYEXP192(temp[0], temp[1], 0x80);
+            break;
+        }
+        case 32:
+        {
+            /* 256 bit key setup */
+            rk[0] = _mm_loadu_si128((const __m128i*) cipherKey);
+            rk[1] = _mm_loadu_si128((const __m128i*) (cipherKey+16));
+            rk[2] = KEYEXP256(rk[0], rk[1], 0x01);
+            rk[3] = KEYEXP256_2(rk[1], rk[2]);
+            rk[4] = KEYEXP256(rk[2], rk[3], 0x02);
+            rk[5] = KEYEXP256_2(rk[3], rk[4]);
+            rk[6] = KEYEXP256(rk[4], rk[5], 0x04);
+            rk[7] = KEYEXP256_2(rk[5], rk[6]);
+            rk[8] = KEYEXP256(rk[6], rk[7], 0x08);
+            rk[9] = KEYEXP256_2(rk[7], rk[8]);
+            rk[10] = KEYEXP256(rk[8], rk[9], 0x10);
+            rk[11] = KEYEXP256_2(rk[9], rk[10]);
+            rk[12] = KEYEXP256(rk[10], rk[11], 0x20);
+            rk[13] = KEYEXP256_2(rk[11], rk[12]);
+            rk[14] = KEYEXP256(rk[12], rk[13], 0x40);
+            break;
+        }
+    }
+}
+
+/* Decryption key setup */
+static void aes_key_setup_dec(__m128i dk[], const __m128i ek[], int rounds)
+{
+    dk[rounds] = ek[0];
+    for (int i = 1; i < rounds; ++i) {
+        dk[rounds - i] = _mm_aesimc_si128(ek[i]);
+    }
+    dk[0] = ek[rounds];
+}
+
+static void block_init(block_state* self, unsigned char* key, int keylen)
+{
+	int nr = 0;
+	switch (keylen) {
+	    case 16: nr = 10; break;
+	    case 24: nr = 12; break;
+	    case 32: nr = 14; break;
+        default:
+            PyErr_SetString(PyExc_ValueError,
+                "AES key must be either 16, 24, or 32 bytes long");
+		    return;
+	}
+	self->rounds = nr;
+	aes_key_setup_enc(self->ek, key, keylen);
+    aes_key_setup_dec(self->dk, self->ek, nr);
+}
+
+static void block_encrypt(block_state* self, const u8* in, u8* out)
+{
+    __m128i m = _mm_loadu_si128((const __m128i*) in);
+    /* first 9 rounds */
+    m = _mm_xor_si128(m, self->ek[0]);
+    m = _mm_aesenc_si128(m, self->ek[1]);
+    m = _mm_aesenc_si128(m, self->ek[2]);
+    m = _mm_aesenc_si128(m, self->ek[3]);
+    m = _mm_aesenc_si128(m, self->ek[4]);
+    m = _mm_aesenc_si128(m, self->ek[5]);
+    m = _mm_aesenc_si128(m, self->ek[6]);
+    m = _mm_aesenc_si128(m, self->ek[7]);
+    m = _mm_aesenc_si128(m, self->ek[8]);
+    m = _mm_aesenc_si128(m, self->ek[9]);
+    if (self->rounds != 10) {
+        /* two additional rounds for AES-192/256 */
+        m = _mm_aesenc_si128(m, self->ek[10]);
+        m = _mm_aesenc_si128(m, self->ek[11]);
+        if (self->rounds == 14) {
+            /* another two additional rounds for AES-256 */
+            m = _mm_aesenc_si128(m, self->ek[12]);
+            m = _mm_aesenc_si128(m, self->ek[13]);
+        }
+    }
+    m = _mm_aesenclast_si128(m, self->ek[self->rounds]);
+    _mm_storeu_si128((__m128i*) out, m);
+}
+
+static void block_decrypt(block_state* self, const u8* in, u8* out)
+{
+	__m128i m = _mm_loadu_si128((const __m128i*) in);
+    /* first 9 rounds */
+    m = _mm_xor_si128(m, self->dk[0]);
+    m = _mm_aesdec_si128(m, self->dk[1]);
+    m = _mm_aesdec_si128(m, self->dk[2]);
+    m = _mm_aesdec_si128(m, self->dk[3]);
+    m = _mm_aesdec_si128(m, self->dk[4]);
+    m = _mm_aesdec_si128(m, self->dk[5]);
+    m = _mm_aesdec_si128(m, self->dk[6]);
+    m = _mm_aesdec_si128(m, self->dk[7]);
+    m = _mm_aesdec_si128(m, self->dk[8]);
+    m = _mm_aesdec_si128(m, self->dk[9]);
+    if (self->rounds != 10) {
+        /* two additional rounds for AES-192/256 */
+        m = _mm_aesdec_si128(m, self->dk[10]);
+        m = _mm_aesdec_si128(m, self->dk[11]);
+        if (self->rounds == 14) {
+            /* another two additional rounds for AES-256 */
+            m = _mm_aesdec_si128(m, self->dk[12]);
+            m = _mm_aesdec_si128(m, self->dk[13]);
+        }
+    }
+    m = _mm_aesdeclast_si128(m, self->dk[self->rounds]);
+    _mm_storeu_si128((__m128i*) out, m);
+}
+
+#include "block_template.c"
diff --git a/src/config.h.in b/src/config.h.in
index d89930c..b520f4e 100644
--- a/src/config.h.in
+++ b/src/config.h.in
@@ -1,5 +1,8 @@
 /* src/config.h.in.  Generated from configure.ac by autoheader.  */
 
+/* Define to 1 if you have the <cpuid.h> header file. */
+#undef HAVE_CPUID_H
+
 /* Define to 1 if you have the declaration of `mpz_powm', and to 0 if you
    don't. */
 #undef HAVE_DECL_MPZ_POWM
@@ -20,6 +23,9 @@
 /* Define to 1 if you have the <limits.h> header file. */
 #undef HAVE_LIMITS_H
 
+/* Define if CC supports -maes */
+#undef HAVE_MAES
+
 /* Define to 1 if your system has a GNU libc compatible `malloc' function, and
    to 0 otherwise. */
 #undef HAVE_MALLOC
@@ -63,6 +69,9 @@
 /* Define to 1 if you have the <wchar.h> header file. */
 #undef HAVE_WCHAR_H
 
+/* Define to 1 if you have the <wmmintrin.h> header file. */
+#undef HAVE_WMMINTRIN_H
+
 /* Define to the address where bug reports for this package should be sent. */
 #undef PACKAGE_BUGREPORT
 
diff --git a/src/cpuid.c b/src/cpuid.c
new file mode 100644
index 0000000..e4d8b27
--- /dev/null
+++ b/src/cpuid.c
@@ -0,0 +1,100 @@
+/*
+ *  cpuid.c: check CPU capabilities
+ *
+ * Written in 2013 by Sebastian Ramacher <sebastian@ramacher.at>
+ *
+ * ===================================================================
+ * The contents of this file are dedicated to the public domain.  To
+ * the extent that dedication to the public domain is not available,
+ * everyone is granted a worldwide, perpetual, royalty-free,
+ * non-exclusive license to exercise all rights associated with the
+ * contents of this file for any purpose whatsoever.
+ * No rights are reserved.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
+ * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
+ * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+ * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ * ===================================================================
+ */
+#include "Python.h"
+#include <stdint.h>
+#include "config.h"
+#ifdef HAVE_CPUID_H
+#include <cpuid.h>
+#endif
+
+#include "pycrypto_compat.h"
+
+/*
+ * The have_aes_ni Python function
+ */
+
+static char have_aes_ni__doc__[] =
+"have_aes_ni() -> bool\n"
+"\n"
+"Return whether AES-NI instructions are available.\n";
+
+static PyObject *
+have_aes_ni(PyObject *self, PyObject *args)
+{
+#ifndef HAVE_CPUID_H
+    Py_INCREF(Py_False);
+    return Py_False;
+#else
+    uint32_t eax, ebx, ecx, edx;
+    /* call cpuid to check if AES-NI instructions are available */
+    if (__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
+        if (ecx & bit_AES) {
+            Py_INCREF(Py_True);
+            return Py_True;
+        }
+    }
+    Py_INCREF(Py_False);
+    return Py_False;
+#endif
+}
+
+/*
+ * Module-level method table and module initialization function
+ */
+
+static PyMethodDef cpuid_methods[] = {
+    {"have_aes_ni", have_aes_ni, METH_NOARGS, have_aes_ni__doc__},
+    {NULL, NULL, 0, NULL}   /* end-of-list sentinel value */
+};
+
+#ifdef IS_PY3K
+static struct PyModuleDef moduledef = {
+	PyModuleDef_HEAD_INIT,
+	"cpuid",
+	NULL,
+	-1,
+	cpuid_methods,
+	NULL,
+	NULL,
+	NULL,
+	NULL
+};
+#endif
+
+PyMODINIT_FUNC
+#ifdef IS_PY3K
+PyInit_cpuid(void)
+#else
+initcpuid(void)
+#endif
+{
+    /* Initialize the module */
+#ifdef IS_PY3K
+    return PyModule_Create(&moduledef);
+#else
+    Py_InitModule("cpuid", cpuid_methods);
+#endif
+}
+
+/* vim:set ts=4 sw=4 sts=4 expandtab: */