Coverage for key_util.py: 88%

1import base64 

2from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey 

3from cryptography.hazmat.primitives import serialization 

4import logging 

5 

6class DHKeyPair: 

7 ''' 

8 DHKeyPair class stores a pair of public_key and private_key. 

9 

10 Attributes: 

11 ``private_key`` (string): The Private Key. 

12 

13 ``public_key`` (string): The Public Key. 

14 ''' 

15 

16 def __init__(self, private_key, public_key): 

17 self.private_key = private_key 

18 self.public_key = public_key 

19 

20def generate_key_pair(): 

21 ''' 

22 Generate a Keypair 

23 

24 Returns: 

25 DHKeyPair: private_key and public_key  

26 ''' 

27 

28 keypair= DHKeyPair(None,None) 

29 try: 

30 inst_private_key = X25519PrivateKey.generate() 

31 inst_public_key = inst_private_key.public_key() 

32 

33 bytes_private_key = inst_private_key.private_bytes( 

34 encoding=serialization.Encoding.DER, 

35 format=serialization.PrivateFormat.PKCS8, 

36 encryption_algorithm=serialization.NoEncryption() 

37 ) 

38 

39 bytes_public_key = inst_public_key.public_bytes( 

40 encoding=serialization.Encoding.DER, 

41 format=serialization.PublicFormat.SubjectPublicKeyInfo 

42 ) 

43 private_key = base64.b64encode(bytes_private_key).decode('utf-8') 

44 public_key = base64.b64encode(bytes_public_key).decode('utf-8') 

45 keypair.private_key=private_key 

46 keypair.public_key=public_key 

47 except Exception as e: 

48 logging.exception(e) 

49 

50 return keypair 

51 

52def generate_shared_key(private_key_str, public_key_str): 

53 ''' 

54 Generates a SharedKey. 

55 

56 Parameters: 

57 

58 ``private_key_str`` (string): Private Key of one party. 

59  

60 ``public_key_str`` (string): Public Key of the other party. 

61 

62 Returns: 

63 string: shared_key in base64 encoded string format 

64 ''' 

65 shared_key=None 

66 try: 

67 private_key = serialization.load_der_private_key( 

68 base64.b64decode(private_key_str), 

69 password=None 

70 ) 

71 public_key = serialization.load_der_public_key( 

72 base64.b64decode(public_key_str) 

73 ) 

74 shared_key = private_key.exchange(public_key) 

75 shared_key = base64.b64encode(shared_key).decode('utf-8') 

76 except Exception as e: 

77 logging.exception(e) 

78 

79 return shared_key