88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
import diffie_hellman
|
||
|
|
|
||
|
|
|
||
|
|
class DiffieHellmanTest(unittest.TestCase):
|
||
|
|
|
||
|
|
def test_private_in_range(self):
|
||
|
|
primes = [5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]
|
||
|
|
for i in primes:
|
||
|
|
self.assertTrue(1 < diffie_hellman.private_key(i) < i)
|
||
|
|
|
||
|
|
# Can fail due to randomness, but most likely will not,
|
||
|
|
# due to pseudo-randomness and the large number chosen
|
||
|
|
def test_private_key_randomness(self):
|
||
|
|
p = 2147483647
|
||
|
|
private_keys = []
|
||
|
|
for i in range(5):
|
||
|
|
private_keys.append(diffie_hellman.private_key(p))
|
||
|
|
self.assertEqual(len(list(set(private_keys))), len(private_keys))
|
||
|
|
|
||
|
|
def test_public_key_correct(self):
|
||
|
|
p = 23
|
||
|
|
g = 5
|
||
|
|
private = 6
|
||
|
|
expected = 8
|
||
|
|
|
||
|
|
actual = diffie_hellman.public_key(p, g, private)
|
||
|
|
self.assertEqual(actual, expected)
|
||
|
|
|
||
|
|
def test_secret_key_correct(self):
|
||
|
|
p = 23
|
||
|
|
public = 19
|
||
|
|
private = 6
|
||
|
|
expected = 2
|
||
|
|
|
||
|
|
actual = diffie_hellman.secret(p, public, private)
|
||
|
|
self.assertEqual(actual, expected)
|
||
|
|
|
||
|
|
def test_secret_key_correct_large_nums(self):
|
||
|
|
p = int("""120227323036150778550155526710966921740030662\
|
||
|
|
69457894729842354923526575959371158734103742634711454153\
|
||
|
|
30066288563005527069961435922404533456428692335628867529\
|
||
|
|
30249953227657883929905072620233073626594386072962776144\
|
||
|
|
69143365881426187411323246174903542571280506720291038940\
|
||
|
|
7991986070558964461330091797026762932543""".replace(
|
||
|
|
"\n", "").replace(" ", ""))
|
||
|
|
public = int("""7520544115435791944292554616920871123548\
|
||
|
|
58559049691782063133092992058683123990461493675163366079\
|
||
|
|
66149689640419216591714331722664409474612463910928128055\
|
||
|
|
99415792293044373353565984826436410603792531597409532111\
|
||
|
|
27577117569121441377056137760635413505489115127155125391\
|
||
|
|
86192176020596861210448363099541947258202188""".replace(
|
||
|
|
"\n", "").replace(" ", ""))
|
||
|
|
private = int("""248347939362593293991108130435688850515\
|
||
|
|
37971354473275017926961991904690152151776307586179022004\
|
||
|
|
17377685436170904594686456961202706692908603181062371925\
|
||
|
|
882""".replace("\n", "").replace(" ", ""))
|
||
|
|
expected = int("""70900735223964890815905879227737819348\
|
||
|
|
80851869892044649134650898046120174656773533145582564442\
|
||
|
|
98779465564310958207858354973848497783442169812282262526\
|
||
|
|
39932672153547963980483673419756271345828771971984887453\
|
||
|
|
01448857224581986445413661898091472983952358126388674082\
|
||
|
|
1363010486083940557620831348661126601106717071""".replace(
|
||
|
|
"\n", "").replace(" ", ""))
|
||
|
|
|
||
|
|
actual = diffie_hellman.secret(p, public, private)
|
||
|
|
self.assertEqual(actual, expected)
|
||
|
|
|
||
|
|
def test_exchange(self):
|
||
|
|
p = 23
|
||
|
|
g = 5
|
||
|
|
|
||
|
|
privateA = diffie_hellman.private_key(p)
|
||
|
|
privateB = diffie_hellman.private_key(p)
|
||
|
|
|
||
|
|
publicA = diffie_hellman.public_key(p, g, privateA)
|
||
|
|
publicB = diffie_hellman.public_key(p, g, privateB)
|
||
|
|
|
||
|
|
secretA = diffie_hellman.secret(p, publicB, privateA)
|
||
|
|
secretB = diffie_hellman.secret(p, publicA, privateB)
|
||
|
|
|
||
|
|
self.assertEqual(secretA, secretB)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|