From 7084955039df6000c708318424b505f948aeb45c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juhani=20Krekel=C3=A4?= <juhani@krekelä.fi>
Date: Thu, 25 Jun 2020 11:40:55 +0300
Subject: [PATCH] Compact ecdh-curve25519

---
 compact_ecdh_curve25519.py | 16 ++++++++++++++++
 test.py                    | 15 +++++++++++++--
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 compact_ecdh_curve25519.py

diff --git a/compact_ecdh_curve25519.py b/compact_ecdh_curve25519.py
new file mode 100644
index 0000000..025f3d6
--- /dev/null
+++ b/compact_ecdh_curve25519.py
@@ -0,0 +1,16 @@
+leu=lambda b:sum(b[i]<<(i*8)for i in range(len(b)))
+p=2**255-19
+decU=lambda b:leu(b[:-1]+bytes([b[-1]&127]))
+mask=lambda x:sum(x>>i|x<<i for i in range(255))
+def cs(s,a,b):d=mask(s)&(a^b);a=a^d;b=b^d;return a,b
+def x25519(k,u):
+	k=bytearray(k);k[0]&=248;k[31]&=127;k[31]|=64;k=leu(k);u=decU(u)
+	x1=u;x2=1;z2=0;x3=u;z3=1;s=0
+	for t in range(254,-1,-1):kt=(k>>t)&1;s^=kt;x2,x3=cs(s,x2,x3);z2,z3=cs(s,z2,z3);s=kt;A=(x2+z2)%p;AA=pow(A,2,p);B=(x2-z2)%p;BB=pow(B,2,p);E=(AA-BB)%p;C=(x3+z3)%p;D=(x3-z3)%p;DA=(D*A)%p;CB=(C*B)%p;x3=pow((DA+CB)%p,2,p);z3=(x1*pow((DA-CB)%p,2,p))%p;x2=(AA*BB)%p;z2=(E*((AA+(121665*E)%p)%p)%p)%p
+	x2,x3=cs(s,x2,x3);z2,z3=cs(s,z2,z3);u=(x2*pow(z2,p-2,p))%p;return bytes(u>>i&255 for i in range(0,255,8))
+pubkey=lambda sk:x25519(sk,b'\x09'+b'\x00'*31)
+def ecdh(pk, sk):
+	k=x25519(sk,pk);o=0
+	for i in k:o|=i
+	if o==0:raise ValueError
+	else:return k
diff --git a/test.py b/test.py
index a80b609..c6fb439 100644
--- a/test.py
+++ b/test.py
@@ -1,5 +1,5 @@
-import secrets
 import compact_chapoly
+import compact_ecdh_curve25519
 
 def test_chacha20_poly1305():
 	test_vectors = [
@@ -13,7 +13,18 @@ def test_chacha20_poly1305():
 		key, nonce, plaintext, ad, ciphertext, tag = map(bytes.fromhex, test_vector)
 		assert compact_chapoly.enc(ad, key, nonce, plaintext) == ciphertext + tag
 		assert compact_chapoly.dec(ad, key, nonce, ciphertext + tag) == plaintext
-		print('Ok')
+		print('chacha20-poly1305: Ok')
+
+def test_ecdh_curve25519():
+	alice_sk = bytes.fromhex('77076d0a7318a57d3c16c17251b26645df4c2f87ebc0992ab177fba51db92c2a')
+	alice_pk = compact_ecdh_curve25519.pubkey(alice_sk)
+	bob_sk = bytes.fromhex('5dab087e624a8a4b79e17f8b83800ee66f3bb1292618b6fd1c2f8b27ff88e0eb')
+	bob_pk = compact_ecdh_curve25519.pubkey(bob_sk)
+	shared_secret = bytes.fromhex('4a5d9d5ba4ce2de1728e3bf480350f25e07e21c947d19e3376f09b3c1e161742')
+	assert compact_ecdh_curve25519.ecdh(alice_pk, bob_sk) == shared_secret
+	assert compact_ecdh_curve25519.ecdh(bob_pk, alice_sk) == shared_secret
+	print('ecdh-curve25519: Ok')
 
 if __name__ == '__main__':
 	test_chacha20_poly1305()
+	test_ecdh_curve25519()