fix to 3DES key localization procedure

6 files changed, 17 insertions, 11 deletions

diff --git a/CHANGES.txt b/CHANGES.txt
index af5c4ddf..427e3f76 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -14,6 +14,7 @@ Github `repo <https://github.com/etingof/pysnmp>`_
 - Fix to SNMP engine boots counter persistence on Python 3.
 - Fix to Pythonized MIBs loading when only .pyc files are
   present (e.g. py2exe/cx_freeze environments).
+- Fix to key localization for 3DES-based encryption.
 
 Revision 4.3.2, released 2016-02-12
 -----------------------------------
diff --git a/pysnmp/proto/secmod/eso/priv/aesbase.py b/pysnmp/proto/secmod/eso/priv/aesbase.py
index f67bdfe0..bd0a19ec 100644
--- a/pysnmp/proto/secmod/eso/priv/aesbase.py
+++ b/pysnmp/proto/secmod/eso/priv/aesbase.py
@@ -28,13 +28,14 @@ class AbstractAes(aes.Aes):
     def localizeKey(self, authProtocol, privKey, snmpEngineID):
         if authProtocol == hmacmd5.HmacMd5.serviceID:
             localPrivKey = localkey.localizeKeyMD5(privKey, snmpEngineID)
-            while ceil(self.keySize // len(localPrivKey)):
+            for count in range(1, int(ceil(self.keySize * 1.0 / len(localPrivKey)))):
                 # noinspection PyDeprecation,PyCallingNonCallable
-                localPrivKey = localPrivKey + md5(localPrivKey).digest()
+                localPrivKey += md5(localPrivKey).digest()
         elif authProtocol == hmacsha.HmacSha.serviceID:
             localPrivKey = localkey.localizeKeySHA(privKey, snmpEngineID)
-            while ceil(self.keySize // len(localPrivKey)):
-                localPrivKey = localPrivKey + sha1(localPrivKey).digest()
+            # RFC mentions this algo generates 480bit key, but only up to 256 bits are used
+            for count in range(1, int(ceil(self.keySize * 1.0 / len(localPrivKey)))):
+                localPrivKey += sha1(localPrivKey).digest()
         else:
             raise error.ProtocolError(
                 'Unknown auth protocol %s' % (authProtocol,)
diff --git a/pysnmp/proto/secmod/eso/priv/des3.py b/pysnmp/proto/secmod/eso/priv/des3.py
index e91f8fed..2388c364 100644
--- a/pysnmp/proto/secmod/eso/priv/des3.py
+++ b/pysnmp/proto/secmod/eso/priv/des3.py
@@ -51,21 +51,22 @@ class Des3(base.AbstractEncryptionService):
                 'Unknown auth protocol %s' % (authProtocol,)
             )
 
+    # 2.1
     def localizeKey(self, authProtocol, privKey, snmpEngineID):
         if authProtocol == hmacmd5.HmacMd5.serviceID:
             localPrivKey = localkey.localizeKeyMD5(privKey, snmpEngineID)
-            while ceil(self.keySize // len(localPrivKey)):
+            for count in range(1, int(ceil(self.keySize * 1.0 / len(localPrivKey)))):
                 # noinspection PyDeprecation,PyCallingNonCallable
-                localPrivKey = localPrivKey + md5(localPrivKey).digest()
+                localPrivKey += localkey.localizeKeyMD5(localPrivKey, snmpEngineID)
         elif authProtocol == hmacsha.HmacSha.serviceID:
             localPrivKey = localkey.localizeKeySHA(privKey, snmpEngineID)
-            while ceil(self.keySize // len(localPrivKey)):
-                localPrivKey = localPrivKey + sha1(localPrivKey).digest()
+            for count in range(1, int(ceil(self.keySize * 1.0 / len(localPrivKey)))):
+                localPrivKey += localkey.localizeKeySHA(localPrivKey, snmpEngineID)
         else:
             raise error.ProtocolError(
                 'Unknown auth protocol %s' % (authProtocol,)
             )
-        return localPrivKey[:self.keySize]  # key+IV
+        return localPrivKey[:self.keySize]
 
     # 5.1.1.1
     def __getEncryptionKey(self, privKey, snmpEngineBoots):
diff --git a/pysnmp/proto/secmod/rfc3414/priv/base.py b/pysnmp/proto/secmod/rfc3414/priv/base.py
index d683ef98..bfd64b6f 100644
--- a/pysnmp/proto/secmod/rfc3414/priv/base.py
+++ b/pysnmp/proto/secmod/rfc3414/priv/base.py
@@ -9,6 +9,7 @@ from pysnmp.proto import error
 
 class AbstractEncryptionService(object):
     serviceID = None
+    keySize = 0
 
     def hashPassphrase(self, authProtocol, privKey):
         raise error.ProtocolError('no encryption')
diff --git a/pysnmp/proto/secmod/rfc3414/priv/des.py b/pysnmp/proto/secmod/rfc3414/priv/des.py
index 16b882bd..4c9ea0b5 100644
--- a/pysnmp/proto/secmod/rfc3414/priv/des.py
+++ b/pysnmp/proto/secmod/rfc3414/priv/des.py
@@ -24,6 +24,8 @@ random.seed()
 
 class Des(base.AbstractEncryptionService):
     serviceID = (1, 3, 6, 1, 6, 3, 10, 1, 2, 2)  # usmDESPrivProtocol
+    keySize = 16
+
     if version_info < (2, 3):
         _localInt = int(random.random() * 0xffffffff)
     else:
@@ -48,7 +50,7 @@ class Des(base.AbstractEncryptionService):
             raise error.ProtocolError(
                 'Unknown auth protocol %s' % (authProtocol,)
             )
-        return localPrivKey[:32]  # key+IV
+        return localPrivKey[:self.keySize]
 
     # 8.1.1.1
     def __getEncryptionKey(self, privKey, snmpEngineBoots):
diff --git a/pysnmp/proto/secmod/rfc3826/priv/aes.py b/pysnmp/proto/secmod/rfc3826/priv/aes.py
index e6541881..1ec0ca63 100644
--- a/pysnmp/proto/secmod/rfc3826/priv/aes.py
+++ b/pysnmp/proto/secmod/rfc3826/priv/aes.py
@@ -83,7 +83,7 @@ class Aes(base.AbstractEncryptionService):
             raise error.ProtocolError(
                 'Unknown auth protocol %s' % (authProtocol,)
             )
-        return localPrivKey[:16]
+        return localPrivKey[:self.keySize]
 
     # 3.2.4.1
     def encryptData(self, encryptKey, privParameters, dataToEncrypt):