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 argparse 

16import string 

17import sys 

18from random import SystemRandom 

19 

20sys.path.append('/opt/xensource/sm/') 

21import util 

22 

23import XenAPI 

24 

25PROGRAM_NAME = 'keymanagerutil' 

26 

27 

28def load_key(key_hash, vdi_uuid): 

29 """ 

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

31 specified key_hash from the key store 

32 """ 

33 _check_key(key_hash, vdi_uuid) 

34 try: 

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

36 return key 

37 except KeyLookUpError: 

38 return None 

39 

40 

41def _check_key(key_hash, vdi_uuid): 

42 with util.ApiSession(PROGRAM_NAME) as session: 

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 

56 

57class InputError(Exception): 

58 def __init__(self, message): 

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

60 

61 

62class KeyLookUpError(Exception): 

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

64 

65 def __init__(self, message): 

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

67 

68 

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

70 """ 

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

72 will be interpreted by XenRT. 

73 """ 

74 data = {} 

75 if key: 

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

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

78 data['key_hash'] = key_hash 

79 print(json.dumps(data)) 

80 

81 

82KEYSTORE_PATH = '/tmp/keystore.json' 

83 

84 

85def _read_keystore(): 

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

87 if os.path.isfile(KEYSTORE_PATH): 

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

89 key_store = json.load(key_store_file) 

90 for key_hash in key_store: 

91 key_base64 = key_store[key_hash] 

92 key = base64.b64decode(key_base64) 

93 key_store[key_hash] = key 

94 return key_store 

95 else: 

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

97 return {} 

98 

99 

100def _write_keystore(key_store): 

101 """ 

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

103 created if it does not exist. 

104 """ 

105 for key_hash in key_store: 

106 key = key_store[key_hash] 

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

108 key_store[key_hash] = key_base64 

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

110 json.dump(key_store, key_store_file) 

111 key_store_file.write("\n") 

112 

113 

114class KeyManager(object): 

115 """ 

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

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

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

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

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

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

122 

123 """ 

124 

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

126 self.key_type = key_type 

127 self.key_length = key_length 

128 self.key = key 

129 self.key_hash = key_hash 

130 

131 def __add_to_keystore(self): 

132 """ 

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

134 key. 

135 """ 

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

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

138 key_store = _read_keystore() 

139 key_store[self.key_hash] = self.key 

140 _write_keystore(key_store) 

141 

142 def __hash_key(self): 

143 

144 # hash the given key - requires key 

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

146 raise InputError("Need key to hash") 

147 

148 hash_it = hashlib.new('sha256') 

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

150 hash_it.update(self.key) 

151 self.key_hash = hash_it.hexdigest() 

152 return self.key_hash 

153 

154 def generate(self): 

155 """ 

156 generate the encryption key 

157 Hash the generated key 

158 Update the key store with key and hash 

159 """ 

160 util.SMlog('Generating key') 

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

162 self.key_hash = self.__hash_key() 

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

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

165 self.__add_to_keystore() 

166 

167 def get_key(self, log_key_info=True): 

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

169 if not self.key_hash: 

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

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

172 

173 key_store = _read_keystore() 

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

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

176 _print_key_info(key=key) 

177 if not key: 

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

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

180 

181 return key 

182 

183 def get_keyhash(self): 

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

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

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

187 key_store = _read_keystore() 

188 try: 

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

190 key_hash = hashes[0] if hashes else None 

191 _print_key_info(key_hash=key_hash) 

192 except ValueError: 

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

194 

195 def update_keystore(self): 

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

197 

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

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

200 

201 key_store = _read_keystore() 

202 if self.key_hash in key_store: 

203 key_store[self.key_hash] = self.key 

204 else: 

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

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

207 _write_keystore(key_store) 

208 

209 

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

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

212 return AlphaNumericKeyGenerator(key_length=key_length) 

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

214 return RandomKeyGenerator(key_length=key_length) 

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

216 return WeakKeyGenerator() 

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

218 return StrongKeyGenerator() 

219 else: 

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

221 " should be specified to generate the key") 

222 

223 

224class RandomKeyGenerator(object): 

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

226 

227 def __init__(self, key_length): 

228 self.key_length = key_length 

229 

230 def generate(self): 

231 """Generate a completely random byte sequence""" 

232 return os.urandom(self.key_length) 

233 

234 

235class StrongKeyGenerator(RandomKeyGenerator): 

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

237 

238 def __init__(self): 

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

240 

241 

242class WeakKeyGenerator(RandomKeyGenerator): 

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

244 

245 def __init__(self): 

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

247 

248 

249class AlphaNumericKeyGenerator(object): 

250 """Generates alphanumeric keys""" 

251 

252 def __init__(self, key_length=None): 

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

254 

255 def generate(self): 

256 """Generate a completely random alphanumeric sequence""" 

257 keys_from = string.ascii_letters + string.digits 

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

259 

260 

261if __name__ == '__main__': 261 ↛ 263line 261 didn't jump to line 263, because the condition on line 261 was never true

262 

263 parser = argparse.ArgumentParser() 

264 

265 parser.add_argument('--generatekey', action='store_true', dest='generate', 

266 default=False, 

267 help="Generates the encryption key based on the given either keytype or keylength") 

268 

269 parser.add_argument('--getkey', action='store_true', dest='get_key', 

270 default=False, help="To get the key from the keystore based on the given key hash") 

271 

272 parser.add_argument('--getkeyhash', action='store_true', dest='get_key_hash', 

273 default=False, help="To get the key hash from the keystore based on the given key") 

274 

275 parser.add_argument('--updatekeystore', action='store_true', dest='update_keystore', 

276 default=False, 

277 help="If needs to update the already existing key in the keystore pass the keyHash and new key") 

278 

279 parser.add_argument('--keytype', action='store', dest='key_type', default=None, 

280 help='Type of the key: values expected weak or strong') 

281 

282 parser.add_argument('--keylength', action='store', default=None, type=int, 

283 dest='key_length', 

284 help='length of the encryption key in byte') 

285 

286 parser.add_argument('--keyhash', action='store', dest='key_hash', default=None, 

287 help='Encryption key') 

288 

289 parser.add_argument('--key', action='store', dest='key', default=None, 

290 help='Base64-encoded encryption key') 

291 

292 parser_input = parser.parse_args() 

293 

294 if parser_input.key: 

295 parser_input.key = base64.b64decode(parser_input.key) 

296 

297 if parser_input.generate: 

298 KeyManager(key_type=parser_input.key_type, key_length=parser_input.key_length).generate() 

299 elif parser_input.get_key: 

300 KeyManager(key_hash=parser_input.key_hash).get_key() 

301 elif parser_input.get_key_hash: 

302 KeyManager(key=parser_input.key).get_keyhash() 

303 elif parser_input.update_keystore: 

304 KeyManager(key_hash=parser_input.key_hash, key=parser_input.key).update_keystore()