pandaharvester/harvestermisc/superfacility_utils.py

0001 import os
0002 import json
0003 import time
0004 import jwt
0005 from authlib.integrations.requests_client import OAuth2Session
0006 from authlib.oauth2.rfc7523 import PrivateKeyJWT
0007
0008
0009 class SuperfacilityClient:
0010     def __init__(self, cred_dir, base_url="https://api.nersc.gov/api/v1.2"):
0011         # cred_dir is a symlink pointing to the latest versioned cred directory
0012         self.cred_dir = os.path.abspath(cred_dir)
0013         self.base_url = base_url.rstrip("/")
0014         self.token_url = "https://oidc.nersc.gov/c2id/token"
0015         self._session = None
0016         self._token = None
0017         self._token_exp = 0
0018
0019     def _load_creds(self):
0020         real_dir = os.path.realpath(self.cred_dir)
0021         id_path = os.path.join(real_dir, "clientid.txt")
0022         key_path = os.path.join(real_dir, "priv_key.jwk")
0023         with open(id_path, "r") as f:
0024             client_id = f.read().strip()
0025         with open(key_path, "r") as f:
0026             private_key = json.load(f)
0027         return client_id, private_key
0028
0029     def _fetch_token(self):
0030         client_id, private_key = self._load_creds()
0031         self._session = OAuth2Session(
0032             client_id,
0033             private_key,
0034             PrivateKeyJWT(self.token_url),
0035             grant_type="client_credentials",
0036             token_endpoint=self.token_url,
0037         )
0038         token_response = self._session.fetch_token()
0039         access_token = token_response.get("access_token")
0040         if access_token is None:
0041             raise RuntimeError("Failed to fetch access_token using OAuth2Session")
0042
0043         decoded = jwt.decode(access_token, options={"verify_signature": False})
0044         exp = decoded.get("exp")
0045         if exp is None:
0046             raise RuntimeError("Fetched token has no 'exp' attribute!.")
0047
0048         self._token = access_token
0049         self._token_exp = int(exp)
0050
0051     def _ensure_token(self):
0052         now = int(time.time())
0053         if self._token is None or (now + 30) >= self._token_exp:
0054             self._fetch_token()
0055
0056     def _request(self, method, path, **kwargs):
0057         self._ensure_token()
0058         url = f"{self.base_url}{path}"
0059         r = self._session.request(method=method, url=url, **kwargs)
0060         r.raise_for_status()
0061         return r
0062
0063     def get(self, path, **kwargs):
0064         return self._request("GET", path, **kwargs)
0065
0066     def post(self, path, **kwargs):
0067         return self._request("POST", path, **kwargs)
0068
0069     def delete(self, path, **kwargs):
0070         return self._request("DELETE", path, **kwargs)