From dbad1f74d1537d0aec9da90322800cdc756ffc0d Mon Sep 17 00:00:00 2001
From: Enol Fernandez <enol.fernandez@egi.eu>
Date: Thu, 7 Jul 2022 07:58:01 +0100
Subject: [PATCH] Add a decorator function for passing the OIDC URL

It mostly duplicates code, but it works.
---
 fedcloudclient/checkin.py    |  8 ++---
 fedcloudclient/decorators.py | 67 ++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/fedcloudclient/checkin.py b/fedcloudclient/checkin.py
index 8bf2099..cced568 100644
--- a/fedcloudclient/checkin.py
+++ b/fedcloudclient/checkin.py
@@ -17,6 +17,7 @@
     DEFAULT_OIDC_URL,
     oidc_access_token_params,
     oidc_params,
+    oidc_params_with_url,
     oidc_refresh_token_params,
 )
 
@@ -319,11 +320,10 @@ def check(oidc_refresh_token, oidc_access_token):
 
 
 @token.command()
-@oidc_params
-def list_vos(access_token):
+@oidc_params_with_url
+def list_vos(access_token, oidc_url):
     """
     List VO membership(s) of access token
     """
-
-    vos = token_list_vos(access_token, DEFAULT_OIDC_URL)
+    vos = token_list_vos(access_token, oidc_url)
     print("\n".join(vos))
diff --git a/fedcloudclient/decorators.py b/fedcloudclient/decorators.py
index 0fe620b..47776b6 100644
--- a/fedcloudclient/decorators.py
+++ b/fedcloudclient/decorators.py
@@ -177,6 +177,73 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def oidc_params_with_url(func):
+    """
+    Decorator for OIDC parameters.
+    Get access token from oidc-* parameters and replace them in the wrapper function
+    Also adds the OIDC URL as part of the call to the inner function
+    """
+
+    @optgroup.group("OIDC token", help="Choose one of options for providing token")
+    @optgroup.option(
+        "--oidc-agent-account",
+        help="Account name in oidc-agent",
+        envvar="OIDC_AGENT_ACCOUNT",
+        metavar="account",
+    )
+    @optgroup.option(
+        "--oidc-access-token",
+        help="OIDC access token",
+        envvar="OIDC_ACCESS_TOKEN",
+        metavar="token",
+    )
+    @optgroup.option(
+        "--oidc-refresh-token",
+        help="OIDC refresh token. Require also client ID and secret",
+        envvar="OIDC_REFRESH_TOKEN",
+        metavar="token",
+    )
+    @optgroup.option(
+        "--oidc-client-id",
+        help="OIDC client ID",
+        envvar="OIDC_CLIENT_ID",
+        metavar="id",
+    )
+    @optgroup.option(
+        "--oidc-client-secret",
+        help="OIDC client secret",
+        envvar="OIDC_CLIENT_SECRET",
+        metavar="secret",
+    )
+    @optgroup.option(
+        "--oidc-url",
+        help="OIDC identity provider URL",
+        envvar="OIDC_URL",
+        default=DEFAULT_OIDC_URL,
+        show_default=True,
+        metavar="provider-url",
+    )
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        from fedcloudclient.checkin import get_access_token
+
+        oidc_url = kwargs.pop("oidc_url")
+
+        access_token = get_access_token(
+            kwargs.pop("oidc_access_token"),
+            kwargs.pop("oidc_refresh_token"),
+            kwargs.pop("oidc_client_id"),
+            kwargs.pop("oidc_client_secret"),
+            oidc_url,
+            kwargs.pop("oidc_agent_account"),
+        )
+        kwargs["access_token"] = access_token
+        kwargs["oidc_url"] = oidc_url
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 def openstack_params(func):
     """
     Decorator for OpenStack authentication parameters