44import pytest
55from binascii import hexlify , unhexlify
66
7- from .curves import NIST192p , NIST224p , NIST256p , NIST384p , NIST521p
7+ try :
8+ import unittest2 as unittest
9+ except ImportError :
10+ import unittest
11+
12+ from .curves import (
13+ NIST192p ,
14+ NIST224p ,
15+ NIST256p ,
16+ NIST384p ,
17+ NIST521p ,
18+ BRAINPOOLP160r1 ,
19+ )
820from .curves import curves
9- from .ecdh import ECDH , InvalidCurveError , InvalidSharedSecretError , NoKeyError
21+ from .ecdh import (
22+ ECDH ,
23+ InvalidCurveError ,
24+ InvalidSharedSecretError ,
25+ NoKeyError ,
26+ NoCurveError ,
27+ )
1028from .keys import SigningKey , VerifyingKey
1129
1230
@@ -26,6 +44,19 @@ def test_ecdh_each(vcurve):
2644 assert secret1 == secret2
2745
2846
47+ def test_ecdh_both_keys_present ():
48+ key1 = SigningKey .generate (BRAINPOOLP160r1 )
49+ key2 = SigningKey .generate (BRAINPOOLP160r1 )
50+
51+ ecdh1 = ECDH (BRAINPOOLP160r1 , key1 , key2 .verifying_key )
52+ ecdh2 = ECDH (private_key = key2 , public_key = key1 .verifying_key )
53+
54+ secret1 = ecdh1 .generate_sharedsecret_bytes ()
55+ secret2 = ecdh2 .generate_sharedsecret_bytes ()
56+
57+ assert secret1 == secret2
58+
59+
2960def test_ecdh_no_public_key ():
3061 ecdh1 = ECDH (curve = NIST192p )
3162
@@ -38,6 +69,44 @@ def test_ecdh_no_public_key():
3869 ecdh1 .generate_sharedsecret_bytes ()
3970
4071
72+ class TestECDH (unittest .TestCase ):
73+ def test_load_key_from_wrong_curve (self ):
74+ ecdh1 = ECDH ()
75+ ecdh1 .set_curve (NIST192p )
76+
77+ key1 = SigningKey .generate (BRAINPOOLP160r1 )
78+
79+ with self .assertRaises (InvalidCurveError ) as e :
80+ ecdh1 .load_private_key (key1 )
81+
82+ self .assertIn ("Curve mismatch" , str (e .exception ))
83+
84+ def test_generate_without_curve (self ):
85+ ecdh1 = ECDH ()
86+
87+ with self .assertRaises (NoCurveError ) as e :
88+ ecdh1 .generate_private_key ()
89+
90+ self .assertIn ("Curve must be set" , str (e .exception ))
91+
92+ def test_load_bytes_without_curve_set (self ):
93+ ecdh1 = ECDH ()
94+
95+ with self .assertRaises (NoCurveError ) as e :
96+ ecdh1 .load_private_key_bytes (b"\x01 " * 32 )
97+
98+ self .assertIn ("Curve must be set" , str (e .exception ))
99+
100+ def test_set_curve_from_received_public_key (self ):
101+ ecdh1 = ECDH ()
102+
103+ key1 = SigningKey .generate (BRAINPOOLP160r1 )
104+
105+ ecdh1 .load_received_public_key (key1 .verifying_key )
106+
107+ self .assertEqual (ecdh1 .curve , BRAINPOOLP160r1 )
108+
109+
41110def test_ecdh_wrong_public_key_curve ():
42111 ecdh1 = ECDH (curve = NIST192p )
43112 ecdh1 .generate_private_key ()
0 commit comments