doma/bin/setup_panda_token.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2021 - 2022
0010
0011 """
0012 Setup Panda IAM token
0013 """
0014
0015 from __future__ import print_function
0016
0017 import argparse
0018 import argcomplete
0019 import base64
0020 import datetime
0021 import json
0022 import os
0023 import socket
0024 import subprocess
0025 import sys
0026
0027 try:
0028     from urllib import urlencode
0029     from urllib2 import urlopen, Request, HTTPError
0030 except ImportError:
0031     from urllib.request import urlopen, Request
0032     from urllib.parse import urlencode
0033     from urllib.error import HTTPError
0034
0035 from pandaclient import panda_api
0036 from pandaclient import Client
0037 # from pandaclient import openidc_utils
0038 from pandaclient import PLogger
0039
0040
0041 def run_command(cmd):
0042     """
0043     Runs a command in an out-of-procees shell.
0044     """
0045     process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, preexec_fn=os.setsid)
0046     stdout, stderr = process.communicate()
0047     if stdout is not None and type(stdout) in [bytes]:
0048         stdout = stdout.decode()
0049     if stderr is not None and type(stderr) in [bytes]:
0050         stderr = stderr.decode()
0051     status = process.returncode
0052     return status, stdout, stderr
0053
0054
0055 def setup_panda_token(verbose=False):
0056     c = panda_api.get_api()
0057     c.hello(verbose)
0058
0059
0060 def get_expire_time(verbose=False):
0061     try:
0062         # token_file = openidc_utils.OpenIdConnect_Utils().get_token_path()
0063         curl = Client._Curl()
0064         curl.verbose = verbose
0065         tmp_log = PLogger.getPandaLogger()
0066         oidc = curl.get_oidc(tmp_log)
0067         token_file = oidc.get_token_path()
0068         if os.path.exists(token_file):
0069             with open(token_file) as f:
0070                 data = json.load(f)
0071                 enc = data['id_token'].split('.')[1]
0072                 enc += '=' * (-len(enc) % 4)
0073                 dec = json.loads(base64.urlsafe_b64decode(enc.encode()))
0074                 exp_time = datetime.datetime.utcfromtimestamp(dec['exp'])
0075                 delta = exp_time - datetime.datetime.utcnow()
0076                 minutes = delta.total_seconds() / 60
0077                 print('Token will expire in %s minutes.' % minutes)
0078                 print('Token expiration time : {0} UTC'.format(exp_time.strftime("%Y-%m-%d %H:%M:%S")))
0079         else:
0080             print("Cannot find token file.")
0081     except Exception as e:
0082         print('failed to decode cached token with {0}'.format(e))
0083
0084
0085 def get_token_info(verbose=False):
0086     # c = panda_api.get_api()
0087     curl = Client._Curl()
0088     curl.verbose = verbose
0089     token_info = curl.get_token_info()
0090     # print(token_info)
0091     if token_info and type(token_info) in [dict]:
0092         for key in token_info:
0093             print("%s: %s" % (key, token_info[key]))
0094         get_expire_time()
0095     else:
0096         print(token_info)
0097
0098
0099 def get_refresh_token_string(verbose=False):
0100     try:
0101         curl = Client._Curl()
0102         curl.verbose = verbose
0103         tmp_log = PLogger.getPandaLogger()
0104         oidc = curl.get_oidc(tmp_log)
0105         token_file = oidc.get_token_path()
0106         if os.path.exists(token_file):
0107             with open(token_file) as f:
0108                 data = json.load(f)
0109                 enc = data['id_token'].split('.')[1]
0110                 enc += '=' * (-len(enc) % 4)
0111                 dec = json.loads(base64.urlsafe_b64decode(enc.encode()))
0112                 exp_time = datetime.datetime.utcfromtimestamp(dec['exp'])
0113                 delta = exp_time - datetime.datetime.utcnow()
0114                 minutes = delta.total_seconds() / 60
0115                 print('Token will expire in %s minutes.' % minutes)
0116                 print('Token expiration time : {0} UTC'.format(exp_time.strftime("%Y-%m-%d %H:%M:%S")))
0117                 if delta < datetime.timedelta(minutes=0):
0118                     print("Token already expired. Cannot refresh.")
0119                     return False, None, None
0120                 return True, data['refresh_token'], delta
0121         else:
0122             print("Cannot find token file.")
0123     except Exception as e:
0124         print('failed to decode cached token with {0}'.format(e))
0125     return False, None, None
0126
0127
0128 def oidc_refresh_token(oidc, token_endpoint, client_id, client_secret, refresh_token_string):
0129     if oidc.verbose:
0130         oidc.log_stream.debug('refreshing token')
0131     data = {'client_id': client_id,
0132             'client_secret': client_secret,
0133             'grant_type': 'refresh_token',
0134             'refresh_token': refresh_token_string}
0135     rdata = urlencode(data).encode()
0136     req = Request(token_endpoint, rdata)
0137     req.add_header('content-type', 'application/x-www-form-urlencoded')
0138     try:
0139         conn = urlopen(req)
0140         text = conn.read()
0141         if oidc.verbose:
0142             oidc.log_stream.debug(text)
0143         id_token = json.loads(text)['id_token']
0144         text = json.dumps(json.loads(text))
0145         with open(oidc.get_token_path(), 'w') as f:
0146             f.write(text)
0147         return True, id_token
0148     except HTTPError as e:
0149         return False, 'code={0}. reason={1}. description={2}'.format(e.code, e.reason, e.read())
0150     except Exception as e:
0151         return False, str(e)
0152
0153
0154 def refresh_token(minutes=30, verbose=False):
0155     curl = Client._Curl()
0156     curl.verbose = verbose
0157     tmp_log = PLogger.getPandaLogger()
0158     oidc = curl.get_oidc(tmp_log)
0159
0160     status, refresh_token, delta = get_refresh_token_string()
0161     if not status:
0162         print("Cannot refresh token.")
0163         return False
0164
0165     print("Fetching auth configuration from: %s" % str(oidc.auth_config_url))
0166     s, o = oidc.fetch_page(oidc.auth_config_url)
0167     if not s:
0168         print("Failed to get Auth configuration: " + o)
0169         return False
0170     auth_config = o
0171
0172     print("Fetching endpoint configuration from: %s" % str(auth_config['oidc_config_url']))
0173     s, o = oidc.fetch_page(auth_config['oidc_config_url'])
0174     if not s:
0175         print("Failed to get endpoint configuration: " + o)
0176         return False
0177     endpoint_config = o
0178
0179     # s, o = oidc.refresh_token(endpoint_config['token_endpoint'], auth_config['client_id'],
0180     #                           auth_config['client_secret'], refresh_token)
0181     s, o = oidc_refresh_token(oidc, endpoint_config['token_endpoint'], auth_config['client_id'],
0182                               auth_config['client_secret'], refresh_token)
0183     if not s:
0184         print("Failed to refresh token: " + o)
0185         if delta < datetime.timedelta(minutes=minutes):
0186             print("The left lifetime of the token is less than required %s minutes" % minutes)
0187             return False
0188         else:
0189             return True
0190     else:
0191         print("Success to refresh token: " + o)
0192         if delta < datetime.timedelta(minutes=minutes):
0193             print("The left lifetime of the token is less than required %s minutes" % minutes)
0194             return False
0195         else:
0196             return True
0197     return True
0198
0199
0200 def refresh_token_action(min_minutes, email, stop_idds, verbose=False):
0201     status = refresh_token(min_minutes, verbose=verbose)
0202     if not status:
0203         if email:
0204             hostname = socket.getfqdn()
0205             cmd = """sendmail "%s"<<EOF
0206                      Subject: iDDS token at %s is going to expire.
0207
0208                      iDDS token is googing to expire. Please refresh it.
0209                      EOF
0210                      """
0211             cmd = cmd % (email, hostname)
0212             status, stdout, stderr = run_command(cmd)
0213             print("Notifying %s: status: %s, stdout: %s, stderr: %s" % (email, status, stdout, stderr))
0214         if stop_idds:
0215             cmd = "supervisorctl start all"
0216             status, stdout, stderr = run_command(cmd)
0217             print("Stop idds %s: status: %s, stdout: %s, stderr: %s" % (cmd, status, stdout, stderr))
0218
0219
0220 def get_parser():
0221     """
0222     Return the argparse parser.
0223     """
0224     oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
0225     # subparsers = oparser.add_subparsers()
0226
0227     oparser.add_argument('--setup', '-s', default=False, action='store_true', help="Setup token.")
0228     oparser.add_argument('--info', '-i', default=False, action='store_true', help="Print token info.")
0229     oparser.add_argument('--refresh', '-r', default=False, action='store_true', help="Refresh token.")
0230     oparser.add_argument('--verbose', '-v', default=False, action='store_true', help="Verbose.")
0231     oparser.add_argument('--panda_config_root', '-p', default=None, action='store', help="Panda config root path. If it's not specified and PANDA_CONFIG_ROOT is not defined, ~/.panda/ will be used.")
0232     oparser.add_argument('--email', '-e', default=None, action='store', help="Email to notify when failing to refresh token and the token will expire.")
0233     oparser.add_argument('--stop_idds', '-t', default=False, action='store_true', help="Stop idds when failing to refresh token and the token will expire.")
0234     oparser.add_argument('--min_minutes', '-m', default=35, action='store', type=int, help="Minimal minutes for a token before email notification and stop idds.")
0235
0236     return oparser
0237
0238
0239 if __name__ == '__main__':
0240     oparser = get_parser()
0241     argcomplete.autocomplete(oparser)
0242
0243     if len(sys.argv) == 1:
0244         oparser.print_help()
0245         sys.exit(-1)
0246
0247     arguments = sys.argv[1:]
0248     args = oparser.parse_args(arguments)
0249     if args.panda_config_root:
0250         os.environ["PANDA_CONFIG_ROOT"] = args.panda_config_root
0251     else:
0252         if "PANDA_CONFIG_ROOT" not in os.environ:
0253             os.environ["PANDA_CONFIG_ROOT"] = "~/.panda"
0254     if not os.path.exists(os.environ["PANDA_CONFIG_ROOT"]):
0255         os.makedirs(os.environ["PANDA_CONFIG_ROOT"])
0256
0257     if args.setup:
0258         setup_panda_token(args.verbose)
0259     elif args.info:
0260         get_token_info(args.verbose)
0261     elif args.refresh:
0262         refresh_token_action(min_minutes=args.min_minutes, email=args.email, stop_idds=args.stop_idds, verbose=args.verbose)
0263     else:
0264         oparser.print_help()
0265         sys.exit(-1)