fix: use write_credentials() from diracx

aldbr · Sep 18, 2023 · 7c6372e · 7c6372e
1 parent 21634fb
commit 7c6372e
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 25 deletions.
diff --git a/src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py b/src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py
@@ -440,7 +440,6 @@ def export_exchangeProxyForToken(self):
                 TokenResponse(
                     access_token=create_token(payload, authSettings),
                     expires_in=authSettings.access_token_expire_minutes * 60,
-                    state="None",
                 ).dict()
             )
         except Exception as e:

diff --git a/src/DIRAC/FrameworkSystem/scripts/dirac_login.py b/src/DIRAC/FrameworkSystem/scripts/dirac_login.py
@@ -315,23 +315,21 @@ def loginWithCertificate(self):
 
             # Get a token for use with diracx
             if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
-                from diracx.cli import EXPIRES_GRACE_SECONDS  # pylint: disable=import-error
-                from diracx.cli.utils import CREDENTIALS_PATH  # pylint: disable=import-error
+                from diracx.core.utils import write_credentials  # pylint: disable=import-error
+                from diracx.core.models import TokenResponse  # pylint: disable=import-error
 
                 res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
                 if not res["OK"]:
                     return res
-                CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True)
-                expires = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(
-                    seconds=res["Value"]["expires_in"] - EXPIRES_GRACE_SECONDS
+                token_content = res["Value"]
+                write_credentials(
+                    TokenResponse(
+                        access_token=token_content["access_token"],
+                        expires_in=token_content["expires_in"],
+                        token_type=token_content.get("token_type"),
+                        refresh_token=token_content.get("refresh_token"),
+                    )
                 )
-                credential_data = {
-                    "access_token": res["Value"]["access_token"],
-                    # TODO: "refresh_token":
-                    # TODO: "refresh_token_expires":
-                    "expires": expires.isoformat(),
-                }
-                CREDENTIALS_PATH.write_text(json.dumps(credential_data))
 
         return S_OK()
 

diff --git a/src/DIRAC/FrameworkSystem/scripts/dirac_proxy_init.py b/src/DIRAC/FrameworkSystem/scripts/dirac_proxy_init.py
@@ -241,23 +241,21 @@ def doTheMagic(self):
                     return resultProxyUpload
 
         if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
-            from diracx.cli import EXPIRES_GRACE_SECONDS  # pylint: disable=import-error
-            from diracx.cli.utils import CREDENTIALS_PATH  # pylint: disable=import-error
+            from diracx.core.utils import write_credentials  # pylint: disable=import-error
+            from diracx.core.models import TokenResponse  # pylint: disable=import-error
 
             res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
             if not res["OK"]:
                 return res
-            CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True)
-            expires = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(
-                seconds=res["Value"]["expires_in"] - EXPIRES_GRACE_SECONDS
+            token_content = res["Value"]
+            write_credentials(
+                TokenResponse(
+                    access_token=token_content["access_token"],
+                    expires_in=token_content["expires_in"],
+                    token_type=token_content.get("token_type"),
+                    refresh_token=token_content.get("refresh_token"),
+                )
             )
-            credential_data = {
-                "access_token": res["Value"]["access_token"],
-                # TODO: "refresh_token":
-                # TODO: "refresh_token_expires":
-                "expires": expires.isoformat(),
-            }
-            CREDENTIALS_PATH.write_text(json.dumps(credential_data))
 
         return S_OK()