Coverage for libs/sm/plugins/keymanagerutil.py: 71%

1#!/usr/bin/python3 

2 

3""" 

4Helper module used for XenRT testing of the VDI encryption feature (REQ-718). 

5This module implements the key lookup plugin interface, so if it is 

6installed, SM will use it to retrieve keys based on their hashes. 

7This key store is backed by a file stored on disk in dom0, and helper 

8functions are provided to manipulate it. 

9""" 

10import base64 

11import os 

12import os.path 

13import hashlib 

14import json 

15import string 

16import sys 

17from random import SystemRandom 

18 

19from sm.core import util 

20import XenAPI # pylint: disable=import-error 

21 

22 

23PROGRAM_NAME = 'keymanagerutil' 

24 

25 

26def load_key(key_hash, vdi_uuid): 

27 """ 

28 load_key is called by SM plugin when it needs to find the key for 

29 specified key_hash from the key store 

30 """ 

31 _check_key(key_hash, vdi_uuid) 

32 try: 

33 key = KeyManager(key_hash=key_hash).get_key(log_key_info=False) 

34 return key 

35 except KeyLookUpError: 

36 return None 

37 

38 

39def _check_key(key_hash, vdi_uuid): 

40 session = XenAPI.xapi_local() 

41 session.xenapi.login_with_password('root', '', '', PROGRAM_NAME) 

42 try: 

43 vdi = session.xenapi.VDI.get_by_uuid(vdi_uuid) 

44 sm_config = session.xenapi.VDI.get_sm_config(vdi) 

45 if 'key_hash' in sm_config: 

46 if key_hash != sm_config['key_hash']: 

47 raise Exception('A key was requested with key hash {}' 

48 ' for VDI {}, but it has a different' 

49 ' key_hash in its sm_config:' 

50 ' {}'.format(key_hash, vdi_uuid, sm_config['key_hash'])) 

51 else: 

52 raise Exception('Encryption key requested for VDI {}' 

53 ' whose sm_config does not contain the key_hash' 

54 ' entry. Its sm_config is {}'.format(vdi_uuid, sm_config)) 

55 finally: 

56 session.xenapi.logout() 

57 

58 

59class InputError(Exception): 

60 def __init__(self, message): 

61 super(InputError, self).__init__(message) 

62 

63 

64class KeyLookUpError(Exception): 

65 """Raised when the key / key hash we've requested is not in the keystore""" 

66 

67 def __init__(self, message): 

68 super(KeyLookUpError, self).__init__(message) 

69 

70 

71def _print_key_info(key=None, key_hash=None): 

72 """ 

73 Output the key details as JSON to the standard output. This output 

74 will be interpreted by XenRT. 

75 """ 

76 data = {} 

77 if key: 

78 data['key_base64'] = base64.b64encode(key).decode() 

79 if key_hash: 79 ↛ 81line 79 didn't jump to line 81, because the condition on line 79 was never false

80 data['key_hash'] = key_hash 

81 print(json.dumps(data)) 

82 

83 

84KEYSTORE_PATH = '/tmp/keystore.json' 

85 

86 

87def _read_keystore(): 

88 """If the keystore file exists, returns its contents, otherwise returns an empty dictionary.""" 

89 if os.path.isfile(KEYSTORE_PATH): 

90 with open(KEYSTORE_PATH, "r") as key_store_file: 

91 key_store = json.load(key_store_file) 

92 for key_hash in key_store: 

93 key_base64 = key_store[key_hash] 

94 key = base64.b64decode(key_base64) 

95 key_store[key_hash] = key 

96 return key_store 

97 else: 

98 util.SMlog(f'Keystore path {KEYSTORE_PATH} not found') 

99 return {} 

100 

101 

102def _write_keystore(key_store): 

103 """ 

104 Write the given key store contents to the key store file, which will be 

105 created if it does not exist. 

106 """ 

107 for key_hash in key_store: 

108 key = key_store[key_hash] 

109 key_base64 = base64.b64encode(key).decode() 

110 key_store[key_hash] = key_base64 

111 with open(KEYSTORE_PATH, "w+") as key_store_file: 

112 json.dump(key_store, key_store_file) 

113 key_store_file.write("\n") 

114 

115 

116class KeyManager(object): 

117 """ 

118 KeyManager is a python utility tool for generating and managing the keys in the jey store. 

119 One can request KeyManager to generate the keys, passing just the type of 

120 the key - either strong or weak or even the length of the key. 

121 One can request KeyManger to get the key from the key store by passing key_hash. 

122 One can request KeyManager to get the key_hash from the key store by passing encryption key. 

123 KeyManager maintains the keystore(json record) under /tmp/keystore.json. 

124 

125 """ 

126 

127 def __init__(self, key_type=None, key_length=None, key=None, key_hash=None): 

128 self.key_type = key_type 

129 self.key_length = key_length 

130 self.key = key 

131 self.key_hash = key_hash 

132 

133 def __add_to_keystore(self): 

134 """ 

135 Update the key and key hash in the key store - requires both hash and 

136 key. 

137 """ 

138 if not self.key_hash or not self.key: 138 ↛ 139line 138 didn't jump to line 139, because the condition on line 138 was never true

139 raise InputError("Need both key_hash and key to update into key store") 

140 key_store = _read_keystore() 

141 key_store[self.key_hash] = self.key 

142 _write_keystore(key_store) 

143 

144 def __hash_key(self): 

145 

146 # hash the given key - requires key 

147 if not self.key: 147 ↛ 148line 147 didn't jump to line 148, because the condition on line 147 was never true

148 raise InputError("Need key to hash") 

149 

150 hash_it = hashlib.new('sha256') 

151 hash_it.update(b'\0' * 32) 

152 hash_it.update(self.key) 

153 self.key_hash = hash_it.hexdigest() 

154 return self.key_hash 

155 

156 def generate(self): 

157 """ 

158 generate the encryption key 

159 Hash the generated key 

160 Update the key store with key and hash 

161 """ 

162 util.SMlog('Generating key') 

163 self.key = _get_key_generator(key_length=self.key_length, key_type=self.key_type).generate() 

164 self.key_hash = self.__hash_key() 

165 _print_key_info(key=self.key, key_hash=self.key_hash) 

166 util.SMlog(f'Generated key, hash {self.key_hash}') 

167 self.__add_to_keystore() 

168 

169 def get_key(self, log_key_info=True): 

170 """Fetch the key from the key store based on the key_hash - requires key hash""" 

171 if not self.key_hash: 

172 util.SMlog('No key hash set, cannot retrieve unknown key') 

173 raise InputError("Need key hash to get the key from the key store") 

174 

175 key_store = _read_keystore() 

176 key = key_store.get(self.key_hash, None) 

177 if key and log_key_info: 177 ↛ 178line 177 didn't jump to line 178, because the condition on line 177 was never true

178 _print_key_info(key=key) 

179 if not key: 

180 util.SMlog(f'No key found in keystore for hash {self.key_hash}, known hashes {key_store.keys()}') 

181 raise KeyLookUpError("No keys in the keystore which matches the given key hash") 

182 

183 return key 

184 

185 def get_keyhash(self): 

186 """Fetch the key hash from the key store based on the key - requires key""" 

187 if not self.key: 187 ↛ 188line 187 didn't jump to line 188, because the condition on line 187 was never true

188 raise InputError("Need key to get the key hash from the key store") 

189 key_store = _read_keystore() 

190 try: 

191 hashes = [x for x in key_store.keys() if key_store[x] == self.key] 

192 key_hash = hashes[0] if hashes else None 

193 _print_key_info(key_hash=key_hash) 

194 except ValueError: 

195 raise KeyLookUpError("No key hash in the keystore which matches the given key") 

196 

197 def update_keystore(self): 

198 """If this key hash is already in the key store, update its corresponding key""" 

199 

200 if not (self.key_hash and self.key): 

201 raise InputError("Need key hash and key to update the key store") 

202 

203 key_store = _read_keystore() 

204 if self.key_hash in key_store: 

205 key_store[self.key_hash] = self.key 

206 else: 

207 raise InputError("No existing key in the keystore" 

208 "with key hash {}".format(self.key_hash)) 

209 _write_keystore(key_store) 

210 

211 

212def _get_key_generator(key_length=None, key_type=None): 

213 if key_type == "alphanumeric": 213 ↛ 214line 213 didn't jump to line 214, because the condition on line 213 was never true

214 return AlphaNumericKeyGenerator(key_length=key_length) 

215 elif key_length: 215 ↛ 216line 215 didn't jump to line 216, because the condition on line 215 was never true

216 return RandomKeyGenerator(key_length=key_length) 

217 elif key_type == "weak": 217 ↛ 218line 217 didn't jump to line 218, because the condition on line 217 was never true

218 return WeakKeyGenerator() 

219 elif key_type == "strong": 219 ↛ 222line 219 didn't jump to line 222, because the condition on line 219 was never false

220 return StrongKeyGenerator() 

221 else: 

222 raise InputError("Either key_length in byte or key_type(\"strong OR weak\")" 

223 " should be specified to generate the key") 

224 

225 

226class RandomKeyGenerator(object): 

227 """Generates a completely random key of the specified length""" 

228 

229 def __init__(self, key_length): 

230 self.key_length = key_length 

231 

232 def generate(self): 

233 """Generate a completely random byte sequence""" 

234 return os.urandom(self.key_length) 

235 

236 

237class StrongKeyGenerator(RandomKeyGenerator): 

238 """Generates a completely random 512-bit key""" 

239 

240 def __init__(self): 

241 super(StrongKeyGenerator, self).__init__(key_length=64) 

242 

243 

244class WeakKeyGenerator(RandomKeyGenerator): 

245 """Generates a completely random 256-bit key""" 

246 

247 def __init__(self): 

248 super(WeakKeyGenerator, self).__init__(key_length=32) 

249 

250 

251class AlphaNumericKeyGenerator(object): 

252 """Generates alphanumeric keys""" 

253 

254 def __init__(self, key_length=None): 

255 self.key_length = 64 if key_length is None else key_length 

256 

257 def generate(self): 

258 """Generate a completely random alphanumeric sequence""" 

259 keys_from = string.ascii_letters + string.digits 

260 return ("".join([SystemRandom().choice(keys_from) for _ in range(self.key_length)])).encode("utf-8")