Fixing python implementation

Author: martinthomson

python/http_ece/__init__.py

@@ -8,19 +8,46 @@
 import pyelliptic
 
 keys = {}
+labels = {}
+
+def deriveKey(mode, salt, key=None, dh=None, keyid=None):
+    def buildInfo(base, context):
+        return b"Content-Encoding: " + base + b"\0" + context
+
+    def deriveDH(mode, keyid, dh):
+        def lengthPrefix(key):
+            return struct.pack("!H", len(key)) + key
+
+        if keyid is None:
+            raise Exception(u"'keyid' is not specified with 'dh'")
+        if not keyid in keys:
+            raise Exception(u"'keyid' doesn't identify a key: " + keyid)
+        if not keyid in labels:
+            raise Exception(u"'keyid' doesn't identify a key label: " + keyid)
+        if mode == "encrypt":
+            senderPubKey = keys[keyid].get_pubkey()
+            receiverPubKey = dh
+        elif mode == "decrypt":
+            senderPubKey = dh
+            receiverPubKey = keys[keyid].get_pubkey()
+        else:
+            raise Exception(u"unknown 'mode' specified: " + mode);
+
+        if type(labels[keyid]) == type(u""):
+            labels[keyid] = labels[keyid].encode("utf-8")
+
+        return (keys[keyid].get_ecdh_key(dh),
+                labels[keyid] + b"\0" +
+                    lengthPrefix(receiverPubKey) + lengthPrefix(senderPubKey))
 
-def deriveKey(salt, key=None, dh=None, keyid=None):
     if salt is None or len(salt) != 16:
         raise Exception(u"'salt' must be a 16 octet value")
 
+    context = b''
     if key is not None:
         secret = key
     elif dh is not None:
-        if keyid is None:
-            raise Exception(u"'keyid' is not specified with 'dh'")
-        if keys[keyid] is None:
-            raise Exception(u"'keyid' doesn't identify a key")
-        secret = keys[keyid].get_ecdh_key(dh)
+        (secret, context) = deriveDH(mode=mode, keyid=keyid, dh=dh)
     elif keyid is not None:
         secret = keys[keyid]
     if secret is None:
@@ -30,14 +57,14 @@ def deriveKey(salt, key=None, dh=None, keyid=None):
         algorithm=hashes.SHA256(),
         length=16,
         salt=salt,
-        info=b"Content-Encoding: aesgcm128",
+        info=buildInfo(b"aesgcm128", context),
         backend=default_backend()
     )
     hkdf_nonce = HKDF(
         algorithm=hashes.SHA256(),
         length=12,
         salt=salt,
-        info=b"Content-Encoding: nonce",
+        info=buildInfo(b"nonce", context),
         backend=default_backend()
     )
     return (hkdf_key.derive(secret), hkdf_nonce.derive(secret))
@@ -62,7 +89,8 @@ def decryptRecord(key, nonce, counter, buffer):
         data = data[1+pad:]
         return data
 
-    (key_, nonce_) = deriveKey(salt=salt, key=key, keyid=keyid, dh=dh)
+    (key_, nonce_) = deriveKey(mode="decrypt", salt=salt,
+                               key=key, keyid=keyid, dh=dh)
     if rs < 2:
         raise Exception(u"Record size too small")
     rs += 16 # account for tags
@@ -87,7 +115,8 @@ def encryptRecord(key, nonce, counter, buffer):
         data += encryptor.tag
         return data
 
-    (key_, nonce_) = deriveKey(salt=salt, key=key, keyid=keyid, dh=dh)
+    (key_, nonce_) = deriveKey(mode="encrypt", salt=salt,
+                               key=key, keyid=keyid, dh=dh)
     if rs < 2:
         raise Exception(u"Record size too small")
     rs -= 1 # account for padding

python/setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name='http_ece',
-    version='0.2.0',
+    version='0.3.0',
     author='Martin Thomson',
     author_email='martin.thomson@gmail.com',
     scripts=[],

python/test.py

@@ -14,56 +14,55 @@
 def log(arg):
     if (count == 1):
         print(arg)
+
 def b64e(arg):
     import base64
     return base64.urlsafe_b64encode(arg).decode()
 
 def rlen():
-    return struct.unpack_from('=H', os.urandom(2))[0]
+    return struct.unpack_from("=H", os.urandom(2))[0]
 
 def encryptDecrypt(length, encryptParams, decryptParams=None):
     if decryptParams is None:
         decryptParams = encryptParams
-    log('Salt: ' + b64e(encryptParams['salt']))
+    log("Salt: " + b64e(encryptParams["salt"]))
     input = os.urandom(min(length, maxLen))
-    # input = new Buffer('I am the walrus')
-    log('Input: ' + b64e(input))
-    encrypted = ece.encrypt(input, salt=encryptParams.get('salt'),
-                            key=encryptParams.get('key'),
-                            keyid=encryptParams.get('keyid'),
-                            dh=encryptParams.get('dh'),
-                            rs=encryptParams.get('rs'))
-    log('Encrypted: ' + b64e(encrypted))
-    decrypted = ece.decrypt(encrypted, salt=decryptParams.get('salt'),
-                            key=decryptParams.get('key'),
-                            keyid=decryptParams.get('keyid'),
-                            dh=decryptParams.get('dh'),
-                            rs=decryptParams.get('rs'))
-    log('Decrypted: ' + b64e(decrypted))
+    # input = new Buffer("I am the walrus")
+    log("Input: " + b64e(input))
+    encrypted = ece.encrypt(input, salt=encryptParams.get("salt"),
+                            key=encryptParams.get("key"),
+                            keyid=encryptParams.get("keyid"),
+                            dh=encryptParams.get("dh"),
+                            rs=encryptParams.get("rs"))
+    log("Encrypted: " + b64e(encrypted))
+    decrypted = ece.decrypt(encrypted, salt=decryptParams.get("salt"),
+                            key=decryptParams.get("key"),
+                            keyid=decryptParams.get("keyid"),
+                            dh=decryptParams.get("dh"),
+                            rs=decryptParams.get("rs"))
+    log("Decrypted: " + b64e(decrypted))
     assert input == decrypted
     log("----- OK");
 
-
 def useExplicitKey():
     params = {
-        'key': os.urandom(16),
-        'salt': os.urandom(16),
-        'rs': rlen() + 1
+        "key": os.urandom(16),
+        "salt": os.urandom(16),
+        "rs": rlen() + 1
     }
-    log('Key: ' + b64e(params['key']))
+    log("Key: " + b64e(params["key"]))
     encryptDecrypt(rlen() + 1, params)
 
 
 def exactlyOneRecord():
     length = min(rlen(), maxLen)
     params = {
-        'key': os.urandom(16),
-        'salt': os.urandom(16),
-        'rs': length + 1
+        "key": os.urandom(16),
+        "salt": os.urandom(16),
+        "rs": length + 1
     }
     encryptDecrypt(length, params)
 
-
 def detectTruncation():
     length = min(rlen(), maxLen)
     key = os.urandom(16)
@@ -75,73 +74,70 @@ def detectTruncation():
     try:
         ece.decrypt(encrypted[0:length + 1 + 16], salt=salt, key=key, rs=rs)
     except Exception as e:
-        log('Decryption error: %s' % e.args)
-        log('----- OK')
+        log("Decryption error: %s" % e.args)
+        log("----- OK")
         ok = True
 
     if not ok:
-        raise Exception('Decryption succeeded, but should not have')
-
-
+        raise Exception("Decryption succeeded, but should not have")
 
 def useKeyId():
     keyid = b64e(os.urandom(16))
     key = os.urandom(16)
     ece.keys[keyid] = key
     params = {
-        'keyid': keyid,
-        'salt': os.urandom(16),
-        'rs': rlen() + 1
+        "keyid": keyid,
+        "salt": os.urandom(16),
+        "rs": rlen() + 1
     }
     encryptDecrypt(rlen(), params)
 
-
-# This is a complete crap-shoot: the pre-eminent crypto library in python
-# doesn't even do ECDH; so this doesn't actually work
 def useDH():
     def isUncompressed(k):
         b1 = k.get_pubkey()[0:1]
-        assert struct.unpack("B", b1)[0] == 4, 'is an uncompressed point'
+        assert struct.unpack("B", b1)[0] == 4, "is an uncompressed point"
 
     # the static key is used by the receiver
-    staticKey = pyelliptic.ECC(curve='prime256v1')
+    staticKey = pyelliptic.ECC(curve="prime256v1")
     isUncompressed(staticKey)
     staticKeyId = b64e(staticKey.get_pubkey()[1:])
     ece.keys[staticKeyId] = staticKey
+    ece.labels[staticKeyId] = "P-256"
 
-    log('Receiver private: ' + b64e(staticKey.get_privkey()))
-    log('Receiver public: ' + b64e(staticKey.get_pubkey()))
+    log("Receiver private: " + b64e(staticKey.get_privkey()))
+    log("Receiver public: " + b64e(staticKey.get_pubkey()))
 
     # the ephemeral key is used by the sender
-    ephemeralKey = pyelliptic.ECC(curve='prime256v1')
+    ephemeralKey = pyelliptic.ECC(curve="prime256v1")
     isUncompressed(ephemeralKey)
     ephemeralKeyId = b64e(ephemeralKey.get_pubkey()[1:])
     ece.keys[ephemeralKeyId] = ephemeralKey
+    ece.labels[ephemeralKeyId] = "P-256"
 
-    log('Sender private: ' + b64e(ephemeralKey.get_privkey()))
-    log('Sender public: ' + b64e(ephemeralKey.get_pubkey()))
+    log("Sender private: " + b64e(ephemeralKey.get_privkey()))
+    log("Sender public: " + b64e(ephemeralKey.get_pubkey()))
 
     encryptParams = {
-        'keyid': ephemeralKeyId,
-        'dh': staticKey.get_pubkey(),
-        'salt': os.urandom(16),
-        'rs': rlen() + 1
+        "keyid": ephemeralKeyId,
+        "dh": staticKey.get_pubkey(),
+        "salt": os.urandom(16),
+        "rs": rlen() + 1
     }
     decryptParams = {
-        'keyid': staticKeyId,
-        'dh': ephemeralKey.get_pubkey(),
-        'salt': encryptParams['salt'],
-        'rs': encryptParams['rs']
+        "keyid": staticKeyId,
+        "dh": ephemeralKey.get_pubkey(),
+        "salt": encryptParams["salt"],
+        "rs": encryptParams["rs"]
     }
 
     encryptDecrypt(rlen(), encryptParams, decryptParams)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     for i in list(range(0,count)):
         useExplicitKey()
         exactlyOneRecord()
         detectTruncation()
         useKeyId()
         useDH()
 
-    print('All tests passed.')
+    print("All tests passed.")