Initial AES-NI support

3 files changed, 33 insertions, 3 deletions

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