1# (c) Stefan Countryman 2016-2018 

2 

3""" 

4Safely import GraceDB, registering a warning if the import fails. 

5""" 

6 

7import os 

8import warnings 

9import logging 

10import functools 

11from datetime import datetime 

12from difflib import unified_diff 

13from pathlib import Path 

14from plumbum.cmd import klist, kdestroy, kinit 

15from plumbum import local 

16from plumbum.commands.processes import ProcessExecutionError 

17from llama.com.s3 import PrivateFileCacher 

18from llama.classes import ( 

19 OptionalFeatureWarning, 

20 placeholderclass, 

21 optional_env_var, 

22 MetaClassFactory, 

23) 

24from llama.utils import OUTPUT_DIR 

25 

26LOGGER = logging.getLogger(__name__) 

27GRACEDB_AUTH = optional_env_var( 

28 ['LLAMA_GRACEDB_AUTH'], 

29 errmsg=( 

30 'Set LLAMA_GRACEDB_AUTH to the S3 object key of the LIGO GraceDb ' 

31 'robot keytab to automatically authenticate with ' 

32 'GraceDb. Otherwise, do not authenticate automatically.' 

33 ) 

34)[0] 

35KEYTAB = PrivateFileCacher(GRACEDB_AUTH, 'llama') if GRACEDB_AUTH else None 

36GRACEDB_CLIENT_AUTHENTICATED_METHODS = [ 

37 'writeLog', 

38 'files', 

39 'logs', 

40 'superevent', 

41 'superevents', 

42 'event', 

43 'events', 

44] 

45BASHRC = Path('~/.bashrc').expanduser() 

46CERTDIR = Path(OUTPUT_DIR, "cilogon_cert") 

47CERT = CERTDIR/'CERT_KEY.pem' 

48PRINCIPAL = 'gwhen/robot/gwhen.com@LIGO.ORG' 

49 

50 

51# try loading gracedb if possible 

52try: 

53 from ligo.gracedb.rest import GraceDb as LigoGraceDb, HTTPError 

54except ImportError: 

55 MSG = ("Can't import ligo.gracedb; GraceDB access won't work. " 

56 "Defining placeholder GraceDB vars that will error on use.") 

57 LOGGER.debug(MSG) 

58 warnings.warn(MSG, OptionalFeatureWarning) 

59 LigoGraceDb = placeholderclass('LigoGraceDb', 'ligo.gracedb.rest') 

60 HTTPError = placeholderclass('HTTPError', 'ligo.gracedb.rest', (Exception)) 

61 

62 

63def bashlines(): 

64 """Get the lines in ~/.bashrc""" 

65 return BASHRC.read_text().strip().split('\n') 

66 

67 

68def keytab(): 

69 """Download the Kerberos keytab if it does not exist locally and return the 

70 ``Path`` to that keytab. Raises a ``KeyError`` if ``KEYTAB`` is None (due 

71 to ``LLAMA_GRACEDB_AUTH`` not being set to a LLAMA S3 key for a valid LIGO 

72 robot keytab).""" 

73 if KEYTAB is None: 

74 msg = "env var LLAMA_GRACEDB_AUTH S3 key for LIGO robot keytab not set" 

75 LOGGER.error(msg) 

76 raise KeyError(msg) 

77 key = KEYTAB.get() 

78 LOGGER.info("Keytab file stored locally at: %s", key) 

79 return key 

80 

81 

82def uninstall(): 

83 """Uninstall keytab, run kdestroy, and remove env variable declarations 

84 from .bashrc.""" 

85 res = kdestroy() 

86 LOGGER.info("Destroyed Kerberos creds with kdestroy: %s", res) 

87 cleanedrc = [line for line in bashlines() 

88 if (not line.endswith(' # llama.com.gracedb') and 

89 not line.startswith('export X509_USER_KEY='))] 

90 for var in ('X509_USER_KEY', 'X509_USER_CERT'): 

91 if var in os.environ: 

92 LOGGER.info("Removing credentials from `os.environ[%s]`", var) 

93 del os.environ[var] 

94 newrc = '\n'.join(cleanedrc)+'\n' 

95 diff = unified_diff(bashlines(), cleanedrc, fromfile=str(BASHRC), 

96 tofile='cleaned') 

97 if diff: # pylint: disable=using-constant-test 

98 LOGGER.info("Found autogenerated X509 declarations in %s: %s", BASHRC, 

99 '\n'.join(diff)) 

100 bashbackup = str(BASHRC)+'.orig' 

101 with open(bashbackup, 'w') as backup: 

102 backup.write(BASHRC.read_text()) 

103 LOGGER.info("Backed up old %s to %s", BASHRC, bashbackup) 

104 BASHRC.write_text(newrc) 

105 LOGGER.info("Removed autogenerated X509 env vars from %s", BASHRC) 

106 try: 

107 proxyfile = Path(local['grid-proxy-info']['-path']().strip()) 

108 if proxyfile.exists(): 

109 proxyfile.unlink() 

110 LOGGER.info("Deleted proxy file at %s", proxyfile) 

111 except ProcessExecutionError as err: 

112 LOGGER.info("`grid-proxy-init -file` found no proxy file, err: %s", 

113 err) 

114 if KEYTAB is not None and KEYTAB.localpath.exists(): 

115 KEYTAB.localpath.unlink() 

116 LOGGER.info("Deleted local keytab copy at %s", KEYTAB.localpath) 

117 if (CERT).exists(): 

118 (CERT).unlink() 

119 LOGGER.info("Deleted certificate at %s", CERT) 

120 if CERTDIR.exists(): 

121 CERTDIR.rmdir() 

122 LOGGER.info("Deleted certificate directory %s", CERTDIR) 

123 LOGGER.info("Uninstalled successfully.") 

124 

125 

126def install_keytab(): 

127 """Install the keytab, granting access to GraceDB from this device.""" 

128 res = klist(retcode=(0, 1)) 

129 LOGGER.debug("klist: %s", res) 

130 res = kdestroy.run() 

131 LOGGER.debug("destroyed with kdestroy: %s - %s\n%s", *res) 

132 res = klist.run(retcode=1) 

133 LOGGER.debug("klist (should be empty): %s - %s\n%s", *res) 

134 res = kinit[PRINCIPAL, '-k', '-t', keytab()]() 

135 LOGGER.info("Ran kinit: %s", res) 

136 res = local['ligo-proxy-init']['-k']() 

137 LOGGER.info("Ran ligo-proxy-init: %s", res) 

138 res = local['grid-proxy-info']['-identity']() 

139 LOGGER.info("Ran grid-proxy-info: %s", res) 

140 CERTDIR.mkdir(parents=True, exist_ok=True) 

141 path = Path(local['grid-proxy-info']['-path']().strip()) 

142 LOGGER.debug("Path to new certificate: %s", path) 

143 newpath = CERT 

144 newpath.write_bytes(path.read_bytes()) 

145 LOGGER.info("New certificate placed at: %s", newpath) 

146 LOGGER.info("Updating os.environ with X509_USER_KEY and X509_USER_CERT") 

147 os.environ['X509_USER_KEY'] = str(newpath) 

148 os.environ['X509_USER_CERT'] = str(newpath) 

149 LOGGER.info("Configuring .bashrc...") 

150 fmt = 'export X509_USER_KEY="{}" X509_USER_CERT="{}" # llama.com.gracedb' 

151 envline = fmt.format(newpath, newpath) 

152 if envline not in bashlines(): 

153 LOGGER.info("Adding X509_USER_CERT and X509_USER_KEY to .bashrc") 

154 with BASHRC.open('a') as rcfile: 

155 rcfile.write(envline) 

156 else: 

157 LOGGER.info("X509_USER_CERT and X509_USER_KEY already in .bashrc") 

158 res = klist() 

159 LOGGER.info("klist: %s", res) 

160 

161 

162def gracedb_auth_wrapper(func): 

163 """A wrapper for methods connecting to GraceDb. If ``GRACEDB_AUTH`` is set 

164 to the S3 key for a LIGO robot keytab (controlled by setting environmental 

165 variable ``LLAMA_GRACEDB_AUTH`` to that key), this method will catch 

166 ``RunTime`` and ``ligo.gracedb.exceptions.HTTPError`` errors and then try 

167 to refresh authentication credentials before attempting to call the wrapped 

168 function a second time. 

169 """ 

170 doc_header = f""" 

171 This method will automatically try to install LIGO proxy credentials if 

172 a ``RuntimeError`` or ``ligo.gracedb.exceptions.HTTPError`` 

173 is caught (if ``GRACEDB_AUTH`` is set to the S3 key of a valid LIGO 

174 robot keytab, set by environmental variable ``LLAMA_GRACEDB_AUTH``). 

175 

176 Original docstring: 

177 

178 """ 

179 

180 @functools.wraps(func) 

181 def wrapper(self, *args, **kwargs): 

182 try: 

183 return func(self, *args, **kwargs) 

184 except (HTTPError, RuntimeError) as err: 

185 LOGGER.error("Error caught while running %s with args %s " 

186 "and kwargs %s. Error: %s", func.__name__, args, 

187 kwargs, err) 

188 if not GRACEDB_AUTH: 

189 LOGGER.error("GRACEDB_AUTH not set, not retrying. To " 

190 "automatically try refreshing credentials on an " 

191 "authentication error, set environmental " 

192 "variable LLAMA_GRACEDB_AUTH to a an S3 key to a " 

193 "valid LIGO robot keytab.") 

194 raise 

195 LOGGER.error("LLAMA_GRACEDB_AUTH set, attempting to refresh " 

196 "GraceDb credentials:") 

197 install_keytab() 

198 LOGGER.error("Credentials refreshed, updating client dictionary") 

199 self.__dict__.update( 

200 type(self)( 

201 *args, 

202 fail_if_noauth=True, 

203 _err_on_llama_auth_failure=True, 

204 **kwargs 

205 ).__dict__ 

206 ) 

207 LOGGER.error("Reattempting original call of %s.%s.%s", 

208 type(self).__module__, type(self).__name__, func.__name__) 

209 return func(self, *args, **kwargs) 

210 wrapper.__doc__ = doc_header + func.__doc__ 

211 return wrapper 

212 

213 

214GraceDbAuth = MetaClassFactory(gracedb_auth_wrapper, 

215 meth_names=GRACEDB_CLIENT_AUTHENTICATED_METHODS) 

216 

217 

218class GraceDb(LigoGraceDb, metaclass=GraceDbAuth): 

219 """ 

220 A subclass of ``ligo.gracedb.rest.GraceDb`` that refreshes LLAMA 

221 ``GraceDb`` credentials in the event of an authentication error before 

222 retrying requests. Also checks for installed, non-expired credentials 

223 before initializing and automatically installs LLAMA keytab and refreshes 

224 Kerberos principal if necessary. Only does this if environmental variable 

225 ``LLAMA_GRACEDB_AUTH`` is set to an S3 key for a valid LIGO robot keytab; 

226 otherwise, behaves just like ``ligo.gracedb.rest.GraceDb`` (with some extra 

227 logging on authentication errors). 

228 

229 ``ligo.gracedb.rest.GraceDb`` docstring below: 

230 

231 """ 

232 

233 __doc__ += LigoGraceDb.__doc__ or '' 

234 

235 def __init__(self, *args, **kwargs): 

236 """ 

237 Wraps ``ligo.gracedb.rest.GraceDb.__init__``, setting 

238 ``fail_if_noauth`` to ``True`` by default. If ``fail_if_noauth`` is 

239 ``True`` and X509 certificate credentials are being used, will check if 

240 the X509 certificate has expired and will try to refresh credentials 

241 using ``llama.com.gracedb`` if they have expired. 

242 

243 ``fail_if_noauth`` is set to ``bool(GRACEDB_AUTH)``, which in turn is 

244 ``True`` only if the environmental variable ``LLAMA_GRACEDB_AUTH`` is 

245 set to a non-empty string; if you want this to actually *work*, set 

246 ``LLAMA_GRACEDB_AUTH`` to the S3 key of a LIGO robot keytab for this 

247 server. If ``LLAMA_GRACEDB_AUTH`` is set thus, ``GraceDb`` will 

248 automatically check for a valid Kerberos proxy and will attempt 

249 to refresh it if it detects that it does not exist or is out of date. 

250 

251 Original ``ligo.gracedb.rest.GraceDb.__init__`` docstring: 

252 

253 """ 

254 err_on_auth_failure = kwargs.pop("_err_on_llama_auth_failure", False) 

255 fail_if_noauth = kwargs.pop('fail_if_noauth', True) 

256 # if the initial attempt fails, install robot keytab and retry 

257 try: 

258 super().__init__(*args, fail_if_noauth=fail_if_noauth, **kwargs) 

259 if not fail_if_noauth: 

260 return 

261 # LIGO's GraceDb will already error if the cert files don't exist 

262 if all(c in self.credentials for c in ('cert_file', 'key_file')): 

263 if datetime.now() > self.certificate.not_valid_after: 

264 msg = ("GraceDb credentials expired on " 

265 "GraceDb.__init__. Current expiration date: " + 

266 self.certificate.not_valid_after.isoformat()) 

267 LOGGER.error(msg) 

268 raise RuntimeError(msg) 

269 except RuntimeError as err: 

270 if not (fail_if_noauth and GRACEDB_AUTH): 

271 raise 

272 if err_on_auth_failure: 

273 LOGGER.error("Another auth error despite trying to reinstall " 

274 "LIGO robot keytab. Reraising %s -> %s", 

275 type(err).__name__, err) 

276 raise 

277 LOGGER.error("Error while initializing GraceDb(): %s -> %s", 

278 type(err).__name__, err) 

279 LOGGER.error("Attempting to install GraceDb credentials on the " 

280 "assumption that this was the issue.") 

281 install_keytab() 

282 LOGGER.info("Retrying GraceDb client initialization.") 

283 self.__dict__.update( 

284 type(self)( 

285 *args, 

286 _err_on_llama_auth_failure=True, 

287 fail_if_noauth=fail_if_noauth, 

288 **kwargs 

289 ).__dict__ 

290 ) 

291 return 

292 __init__.__doc__ += LigoGraceDb.__init__.__doc__ or ''