1import logging 

2from base64 import b64decode, b64encode 

3from io import BytesIO 

4 

5import pgpy 

6from hvac.exceptions import InvalidPath 

7from pgpy.constants import (CompressionAlgorithm, HashAlgorithm, KeyFlags, 

8 PubKeyAlgorithm, SymmetricKeyAlgorithm) 

9from sdc_etl_libs.pgp_helpers.pgp import PGP 

10 

11 

12class LocalPGP(PGP): 

13 

14 def __init__(self): 

15 """ 

16 Initialized empty key objects and establishes connection to Vault 

17 """ 

18 super().__init__() 

19 self.public_key = None 

20 self.private_key = None 

21 

22 def __str__(self): 

23 """ 

24 :return: String representation of key pair 

25 """ 

26 return str(self.public_key) + '\n\n' + str(self.private_key) 

27 

28 def set_public_key(self, key_): 

29 """ 

30 :param key_: public key string 

31 :return: None 

32 """ 

33 self.public_key = pgpy.PGPKey() 

34 self.public_key.parse(key_) 

35 

36 def set_private_key(self, key_): 

37 """ 

38 :param key_: private key string 

39 :return: None 

40 """ 

41 self.private_key = pgpy.PGPKey() 

42 self.private_key.parse(key_) 

43 

44 def create_keys(self): 

45 """ 

46 This method creates a public/private key pair object to be used for encryption and decryption. 

47 :return: None 

48 """ 

49 key = pgpy.PGPKey.new(PubKeyAlgorithm.RSAEncryptOrSign, 4096) 

50 uid = pgpy.PGPUID.new('sdc-data-engineering') 

51 key.add_uid( 

52 uid, 

53 usage={KeyFlags.Sign, KeyFlags.EncryptCommunications, KeyFlags.EncryptStorage}, 

54 hashes=[HashAlgorithm.SHA256, HashAlgorithm.SHA384, HashAlgorithm.SHA512, HashAlgorithm.SHA224], 

55 ciphers=[SymmetricKeyAlgorithm.AES256, SymmetricKeyAlgorithm.AES192, SymmetricKeyAlgorithm.AES128], 

56 compression=[ 

57 CompressionAlgorithm.ZLIB, CompressionAlgorithm.BZ2, CompressionAlgorithm.ZIP, 

58 CompressionAlgorithm.Uncompressed 

59 ]) 

60 self.set_public_key(str(key.pubkey)) 

61 self.set_private_key(str(key)) 

62 

63 def encrypt(self, data_): 

64 """ 

65 Performs encryption of data 

66 :param data_: data intended for encryption 

67 :return: encrypted data 

68 """ 

69 if not self.public_key: 

70 msg = 'Attempting to encrypt data without public key' 

71 logging.error(msg) 

72 raise ValueError(msg) 

73 

74 message_ = pgpy.PGPMessage.new(data_.read()) 

75 message_ |= self.private_key.sign(message_) 

76 return BytesIO(bytes(self.public_key.encrypt(message_).__str__(), encoding='utf8')) 

77 

78 def decrypt(self, data_): 

79 """ 

80 Performs decryption of data 

81 :param data_: data intended for decryption 

82 :return: decrypted data 

83 :raises: 

84 ValueError: No private available 

85 TypeError: No decryption attempts worked 

86 """ 

87 if not self.private_key: 

88 msg = 'Attempting to encrypt data without private key' 

89 logging.error(msg) 

90 raise ValueError(msg) 

91 

92 message_ = pgpy.PGPMessage.from_blob(data_.read()) 

93 try: 

94 return BytesIO(bytes(self.private_key.decrypt(message_).message, encoding='utf8')) 

95 except TypeError: 

96 return BytesIO(bytes(self.private_key.decrypt(message_).message)) 

97 

98 def vault_create_or_update_keys(self, path_): 

99 """ 

100 Stores keys within hashicorp vault underthe specified name 

101 :path: path for keys to be stored under 

102 """ 

103 # Extract key as base64 encoded bytes with utf-8 encoding 

104 prepare_key = lambda k: str(b64encode(bytes(str(k), encoding='utf8')), 'utf-8') 

105 

106 self.vault.post_secrets('build-secrets', 'data', path_, { 

107 'public_key': prepare_key(self.public_key), 

108 'private_key': prepare_key(self.private_key) 

109 }) 

110 

111 def vault_get_keys(self, path_): 

112 """ 

113 Retrieves keys from hashicorp vault using specified path 

114 :param: path for keys to be stored under 

115 :return: Dictionary of keys stored in vault 

116 :raises: 

117 InvalidPath: No key was available for use within vault, generate new key pair 

118 KeyError: Unable to set either public or private key from Vault API response 

119 """ 

120 try: 

121 response = self.vault.get_secrets('build-secrets', 'data', path_) 

122 except InvalidPath: 

123 logging.warning('Keys for path {0} does not exist. Creating key pair using path: {0}'.format(path_)) 

124 self.create_keys() 

125 self.vault_create_or_update_keys(path_) 

126 response = self.vault.get_secrets('build-secrets', 'data', path_) 

127 

128 try: 

129 self.set_public_key(b64decode(response['public_key'])) 

130 except KeyError: 

131 logging.warning('Unable to set public key') 

132 try: 

133 self.set_private_key(b64decode(response['private_key'])) 

134 except KeyError: 

135 logging.warning('Unable to set private key')