From 8501f5ab271f451eaa48962bf95cbc2b479098a1 Mon Sep 17 00:00:00 2001
From: John-Mark Gurney <jmg@funkthat.com>
Date: Mon, 31 Jan 2022 12:23:34 -0800
Subject: [PATCH] implement an x25519 class and add various tests for it

---
 lora_comms.py | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 83 insertions(+), 1 deletion(-)

diff --git a/lora_comms.py b/lora_comms.py
index 4f8e03b..10e291c 100644
--- a/lora_comms.py
+++ b/lora_comms.py
@@ -25,6 +25,7 @@
 import os
 import unittest
 
+from binascii import a2b_hex
 from ctypes import Structure, POINTER, CFUNCTYPE, pointer, sizeof
 from ctypes import c_uint8, c_uint16, c_ssize_t, c_size_t, c_uint64, c_int
 from ctypes import CDLL
@@ -149,6 +150,55 @@ def x25519_base(scalar, clamp):
 
 	return bytes(out)
 
+class X25519:
+	'''Class to wrap the x25519 functions into something a bit more
+	usable.  This provides better key ingestion and better support
+	for other key formats.
+
+	Use either the gen method to generate a random key, or the frombytes
+	method.
+
+	a = X25519.gen()
+	b = X25519.gen()
+
+	a.dh(b.getpub()) == b.dh(a.getpub())
+
+	That is, each party generates a key, sends their public part to the
+	other party, and then uses their received public part as an argument
+	to the dh method.  The resulting value will be shared between the
+	two parties.
+	'''
+
+	def __init__(self, key):
+		self.privkey = key
+		self.pubkey = x25519_base(key, 1)
+
+	def dh(self, pub):
+		'''Perform a DH operation using the public part pub.'''
+
+		return x25519_wrap(self.pubkey, self.privkey, pub, 1)
+
+	def getpub(self):
+		'''Get the public part of the key.  This is to be sent
+		to the other party for key exchange.'''
+
+		return self.pubkey
+
+	def getpriv(self):
+		return self.privkey
+
+	@classmethod
+	def gen(cls):
+		'''Generate a random X25519 key.'''
+
+		return cls(x25519_genkey())
+
+	@classmethod
+	def frombytes(cls, key):
+		'''Generate an X25519 key from 32 bytes.'''
+
+		return cls(key)
+
 def comms_process_wrap(state, input):
 	'''A wrapper around comms_process that converts the argument
 	into the buffer, and the returns the message as a bytes string.
@@ -164,7 +214,39 @@ def comms_process_wrap(state, input):
 	return outbuf._from()
 
 class TestX25519(unittest.TestCase):
-	def test_basic(self):
+	PUBLIC_BYTES = EC_PUBLIC_BYTES
+	PRIVATE_BYTES = EC_PRIVATE_BYTES
+
+	def test_class(self):
+		key = X25519.gen()
+
+		pubkey = key.getpub()
+		privkey = key.getpriv()
+
+		apubkey = x25519_base(privkey, 1)
+
+		self.assertEqual(apubkey, pubkey)
+		self.assertEqual(X25519.frombytes(privkey).getpub(), pubkey)
+
+		with self.assertRaises(ValueError):
+			X25519(b'0'*31)
+
+	def test_rfc7748_6_1(self):
+		# KAT from https://datatracker.ietf.org/doc/html/rfc7748#section-6.1
+		apriv = a2b_hex('77076d0a7318a57d3c16c17251b26645df4c2f87ebc0992ab177fba51db92c2a')
+
+		akey = X25519(apriv)
+		self.assertEqual(akey.getpub(), a2b_hex('8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a'))
+
+		bpriv = a2b_hex('5dab087e624a8a4b79e17f8b83800ee66f3bb1292618b6fd1c2f8b27ff88e0eb')
+		bkey = X25519(bpriv)
+		self.assertEqual(bkey.getpub(), a2b_hex('de9edb7d7b7dc1b4d35b61c2ece435373f8343c85b78674dadfc7e146f882b4f'))
+
+		ss = a2b_hex('4a5d9d5ba4ce2de1728e3bf480350f25e07e21c947d19e3376f09b3c1e161742')
+		self.assertEqual(akey.dh(bkey.getpub()), ss)
+		self.assertEqual(bkey.dh(akey.getpub()), ss)
+
+	def test_basic_ops(self):
 		aprivkey = x25519_genkey()
 		apubkey = x25519_base(aprivkey, 1)