fix(auth): updates based on comments

Author: Matt-Carre

scripts/checkyamlcompliance.py

@@ -26,6 +26,10 @@
         match yamldata["kind"]:
             case "Workflow":
                 clst_tmpt = True
+            case "WorkflowTemplate":
+                clst_tmpt = True
+            case "ClusterWorkflowTemplate":
+                clst_tmpt = True
             case _:
                 clst_tmpt = False
     else:

scripts/lintyaml.sh

@@ -1,9 +1,11 @@
 #!/bin/bash
 cd src/python_interface_to_workflows/templates
-for file in *.yaml;
+for file in *.txt;
 do
+mv $file.txt $file.yaml
 argo lint "$file" --offline
 SUCCESSFULLINT=$?
 [ $SUCCESSFULLINT -ne 0 ] && exit 1
+mv $file.yaml $file.txt
 done
 exit 0

src/copier_template/scripts/runthefiles.sh.jinja

@@ -1,5 +1,5 @@
 #!/bin/bash
-cd src/python_interface_to_workflows/workflow_definitions
+cd src/{{project_name}}/workflow_definitions
 for file in *
 do
 uv run "$file"

src/copier_template/src/README.md.jinja

@@ -4,7 +4,7 @@
 
 HOST=https://argo-workflows.workflows.diamond.ac.uk/ (to submit to the production cluster)
 DEFAULT_IMAGE= (usually python 3.10)
-NAMESPACE= (the cluster you wish to run the templates on)
+VISIT= (the Visit you wish to run the template on)
 TOKEN=
 EXPIRY=
 AUTH=

src/copier_template/src/{{ project_name }}/auth/keycloak_checker.py.jinja

@@ -1,10 +1,20 @@
 import os
+from typing import TypedDict, cast
 
 import dotenv
 from keycloak import KeycloakOpenID
 from keycloak.pkce_utils import generate_code_challenge, generate_code_verifier
 
-from {{project_name}}.auth.open_auth_url import open_auth_url
+from {{project_name}}.auth.open_auth_url import token_expired
+
+
+class TokenResponse(TypedDict):
+    access_token: str
+    refresh_token: str
+
+
+class DecodedToken(TypedDict):
+    exp: int
 
 
 def set_token_env_variable() -> str:
@@ -25,26 +35,38 @@ def set_token_env_variable() -> str:
         code_challenge=code_challenge,
         code_challenge_method=code_challenge_method,
     )
-    match open_auth_url(auth_url, port):
-        case True:
-            token: dict[str, str] = (  # pyright: ignore[reportUnknownVariableType]
-                keycloak_openid.token(  # pyright: ignore[reportUnknownMemberType]
-                    grant_type="authorization_code",
-                    code=os.environ.get("AUTH"),  # pyright: ignore
-                    redirect_uri=f"http://localhost:{port}/",
-                    code_verifier=code_verifier,
-                )
-            )
-        case False:
-            token: dict[str, str] = keycloak_openid.refresh_token(  # pyright: ignore
-                str(os.environ.get("REFRESHTOKEN"))  # pyright: ignore
-            )
-    token_info: dict[str, int | str | dict[str, list[str]]] = (  # pyright: ignore
-        keycloak_openid.decode_token(token["access_token"])  # pyright: ignore
+    if token_expired(auth_url, port):
+        token = cast(
+            TokenResponse,
+            keycloak_openid.token(  # pyright: ignore[reportUnknownMemberType]
+                grant_type="authorization_code",
+                code=str(os.environ.get("AUTH")),
+                redirect_uri=f"http://localhost:{port}/",
+                code_verifier=code_verifier,
+            ),
+        )
+    else:
+        token = cast(
+            TokenResponse,
+            keycloak_openid.refresh_token(  # pyright: ignore[reportUnknownMemberType]
+                str(os.environ.get("REFRESHTOKEN"))
+            ),
+        )
+    token_info = cast(
+        DecodedToken,
+        keycloak_openid.decode_token(  # pyright: ignore[reportUnknownMemberType]
+            token["access_token"]
+        ),
     )
-    expire_time = int(token_info["exp"]) + 1500  # pyright: ignore
-    dotenv.set_key("src/.env", "EXPIRY", str(expire_time).strip("'"))
-    dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'"))
-    dotenv.set_key("src/.env", "REFRESHTOKEN", token["refresh_token"].strip("'"))
-    dotenv.load_dotenv(dotenv_path="src/.env", override=True)
-    return token["access_token"]  # pyright: ignore[reportUnknownArgumentType]
+    try:
+        expire_time = int(token_info["exp"]) + 1800
+        dotenv.set_key("src/.env", "EXPIRY", str(expire_time))
+        dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'"))
+        dotenv.set_key("src/.env", "REFRESHTOKEN", token["refresh_token"].strip("'"))
+    except AttributeError:
+        print("ERROR:")
+        exit(1)
+    finally:
+        dotenv.load_dotenv(dotenv_path="src/.env", override=True)
+
+    return token["access_token"]

src/copier_template/src/{{ project_name }}/auth/open_auth_url.py

@@ -29,23 +29,27 @@ def do_GET(self):
             self.wfile.write(b"Missing authorization code.")
 
 
-def open_auth_url(auth_url: str, port: int) -> bool:
+def token_expired(auth_url: str, port: int) -> bool:
     dotenv.load_dotenv(dotenv_path="src/.env", override=True)
-    expiry_str: str = os.environ.get("EXPIRY").strip("'")  # pyright: ignore
+    expiry_str: str = str(os.environ.get("EXPIRY")).strip("'")
     if (expiry_str == "" or int(expiry_str)) <= float(time.time()):
-        httpd = _ReusingHTTPServer(("localhost", port), CallbackHandler)
-        webbrowser.open(auth_url)
-        try:
-            httpd.handle_request()
-            os.environ["AUTH"] = httpd.auth_code
-            dotenv.set_key("src/.env", "AUTH", httpd.auth_code)
-        except OSError:
-            os.environ["AUTH"] = ""
-            print("ERROR: Port in use. Please restart your terminal.")
-            exit(1)
-        finally:
-            httpd.socket.shutdown(socket.SHUT_RDWR)
-            httpd.server_close()
+        _open_auth_url(auth_url, port)
         return True
     else:
         return False
+
+
+def _open_auth_url(auth_url: str, port: int) -> None:
+    httpd = _ReusingHTTPServer(("localhost", port), CallbackHandler)
+    webbrowser.open(auth_url)
+    try:
+        httpd.handle_request()
+        os.environ["AUTH"] = httpd.auth_code
+        dotenv.set_key("src/.env", "AUTH", httpd.auth_code)
+    except OSError:
+        os.environ["AUTH"] = ""
+        print("ERROR: Port in use. Please restart your terminal.")
+        exit(1)
+    finally:
+        httpd.socket.shutdown(socket.SHUT_RDWR)
+        httpd.server_close()

src/copier_template/src/{{ project_name }}/submit_to_argo.py.jinja

@@ -12,7 +12,7 @@ from {{project_name}}.auth.keycloak_checker import set_token_env_variable
 def submit_workflow_to_argo(w: Workflow):
     set_token_env_variable()
     dotenv.load_dotenv(dotenv_path="src/.env", override=True)
-    w.namespace = os.environ.get("NAMESPACE")
+    w.namespace = os.environ.get("VISIT")
     w.workflows_service = WorkflowsService(
         host=str(os.environ.get("HOST")).strip("'"),
         token=str(os.environ.get("TOKEN")).strip("'"),

src/copier_template/src/{{ project_name }}/templates/example.txt.jinja

@@ -1,6 +1,6 @@
 {% raw %}
 apiVersion: argoproj.io/v1alpha1
-kind: Workflow
+kind: WorkflowTemplate
 metadata:
   generateName: hera-example-
   annotations:

src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template.py.jinja

@@ -148,7 +148,7 @@ with Workflow(
     generate_name="hera-example-",  # when running on graphql this should be name
     entrypoint="workflowentry",
     api_version="argoproj.io/v1alpha1",
-    kind="Workflow",  # ClusterWorkflowTemplate", when on graphql
+    kind="WorkflowTemplate",  # ClusterWorkflowTemplate", when on graphql
     labels={"workflows.diamond.ac.uk/science-group-examples": "true"},
     annotations={
         "workflows.argoproj.io/title": "example remade via hera",

src/copier_template/tests/test_keycloak_checker.py.jinja

@@ -7,39 +7,36 @@ from {{project_name}}.auth.keycloak_checker import set_token_env_variable
 
 
 @mark.parametrize(
-    "staging,port,return_present",
+    "return_present",
     [
-        (True, 5173, True),
-        (False, 8000, False),
-        (True, 5173, False),
-        (False, 8000, True),
+        (False),
+        (True),
     ],
 )
 @patch("{{project_name}}.auth.keycloak_checker.dotenv.load_dotenv")
 @patch("{{project_name}}.auth.keycloak_checker.dotenv.set_key")
 @patch("{{project_name}}.auth.keycloak_checker.KeycloakOpenID")
 @patch("{{project_name}}.auth.keycloak_checker.generate_code_verifier")
 @patch("{{project_name}}.auth.keycloak_checker.generate_code_challenge")
-@patch("{{project_name}}.auth.keycloak_checker.open_auth_url")
+@patch("{{project_name}}.auth.keycloak_checker.token_expired")
 def test_set_token_env_variable(
-    mock_open_auth_url: MagicMock,
+    mock_token_expired: MagicMock,
     mock_gen_code_challenge: MagicMock,
     mock_gen_code_verifier: MagicMock,
     mock_gen_keycloak_id: MagicMock,
     mock_set_key: MagicMock,
     mock_load_env: MagicMock,
-    staging: bool,
-    port: int,
     return_present: bool,
 ):
+    port=8000
     mock_gen_code_verifier.return_value = "verifier"
     mock_gen_code_challenge.return_value = ("challenge", "S256")
     os.environ["AUTH"] = "auth_url_code"
     os.environ["REFRESHTOKEN"] = "refresh"
 
     keycloak = MagicMock()
     mock_gen_keycloak_id.return_value = keycloak
-    mock_open_auth_url.return_value = return_present
+    mock_token_expired.return_value = return_present
     keycloak.auth_url.return_value = "https://mock.site"
     token = {
         "access_token": "fake_token",
@@ -50,7 +47,7 @@ def test_set_token_env_variable(
     keycloak.decode_token.return_value = {"exp": 123456789}
     assert set_token_env_variable(staging) == "fake_token"
 
-    mock_open_auth_url.assert_called_once_with("https://mock.site", port)
+    mock_token_expired.assert_called_once_with("https://mock.site", port)
     if return_present:
         keycloak.token.assert_called_once_with(
             grant_type="authorization_code",
@@ -64,7 +61,7 @@ def test_set_token_env_variable(
         keycloak.token.assert_not_called()
     mock_set_key.assert_has_calls(
         [
-            call("src/.env", "EXPIRY", str(123456789 + 1500)),
+            call("src/.env", "EXPIRY", str(123456789 + 1800)),
             call("src/.env", "TOKEN", "fake_token"),
             call("src/.env", "REFRESHTOKEN", "fake_refresh"),
         ]
@@ -73,3 +70,51 @@ def test_set_token_env_variable(
         dotenv_path="src/.env",
         override=True,
     )
+
+@patch("{{project_name}}.auth.keycloak_checker.exit")
+@patch("{{project_name}}.auth.keycloak_checker.print")
+@patch("{{project_name}}.auth.keycloak_checker.dotenv.load_dotenv")
+@patch("{{project_name}}.auth.keycloak_checker.dotenv.set_key")
+@patch("{{project_name}}.auth.keycloak_checker.KeycloakOpenID")
+@patch("{{project_name}}.auth.keycloak_checker.generate_code_verifier")
+@patch("{{project_name}}.auth.keycloak_checker.generate_code_challenge")
+@patch("{{project_name}}.auth.keycloak_checker.token_expired")
+def test_set_token_env_variable_attribute_error(
+    mock_token_expired: MagicMock,
+    mock_gen_code_challenge: MagicMock,
+    mock_gen_code_verifier: MagicMock,
+    mock_gen_keycloak_id: MagicMock,
+    mock_set_key: MagicMock,
+    mock_load_env: MagicMock,
+    mock_print: MagicMock,
+    mock_exit: MagicMock,
+):
+    mock_gen_code_verifier.return_value = "verifier"
+    mock_gen_code_challenge.return_value = ("challenge", "S256")
+    mock_token_expired.return_value = True
+
+    os.environ["REFRESHTOKEN"] = "refresh"
+
+    keycloak = MagicMock()
+    mock_gen_keycloak_id.return_value = keycloak
+
+    token = {
+        "access_token": "fake_token",
+        "refresh_token": "fake_refresh",
+    }
+
+    keycloak.refresh_token.return_value = token
+
+    decoded = MagicMock()
+    decoded.__getitem__.side_effect = AttributeError
+    keycloak.decode_token.return_value = decoded
+
+    set_token_env_variable()
+
+    mock_print.assert_called_once_with("ERROR:")
+    mock_exit.assert_called_once_with(1)
+    mock_set_key.assert_not_called()
+    mock_load_env.assert_called_once_with(
+        dotenv_path="src/.env",
+        override=True,
+    )