Merge pull request #16 from nebari-dev/20230531eae

.github/workflows/test.yaml

@@ -0,0 +1,55 @@
+name: "Tests"
+
+on:
+  pull_request:
+    paths:
+      - ".github/workflows/tests.yml"
+      - "nebari_workflow_controller/**"
+  push:
+    branches:
+      - main
+
+jobs:
+  pre-commit:
+    name: "pre-commit"
+    runs-on: ubuntu-latest
+    steps:
+        - name: "Checkout repo"
+          uses: actions/checkout@v3
+          with:
+            fetch-depth: 0
+
+        - name: "Setup Python"
+          uses: actions/setup-python@v3
+
+        - name: "Run pre-commit"
+          uses: pre-commit/[email protected]
+
+  pytest:
+    name: "pytest"
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
+      fail-fast: false
+    steps:
+      - name: "Checkout repo"
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: "Setup Python"
+        uses: actions/setup-python@v3
+
+      - name: "Install package"
+        run: |
+          pip install .[dev]
+
+      - name: "Test package"
+        run: |
+          pytest --version
+          pytest -vvv

.gitignore

@@ -131,3 +131,6 @@ dmypy.json
 
 # vscode
 .vscode/
+
+# Mac
+.DS_Store

README.md

@@ -1,5 +1,5 @@
 # Nebari Workflow Controller
-A kubernetes admission controller to enable volumeMount permisions on Argo Workflows on Nebari and provide a convenience method for deploying jupyterlab-like workflows for users.
+A kubernetes admission controller to enable volumeMount permissions on Argo Workflows on Nebari and provide a convenience method for deploying jupyterlab-like workflows for users.
 
 # Run project
 - `pip install .`

nebari_workflow_controller/utils.py

@@ -1,9 +1,12 @@
 import base64
+import json
 import logging
 import os
+import re
 import traceback
 
 from keycloak import KeycloakAdmin
+from keycloak.exceptions import KeycloakGetError
 from kubernetes import client, config
 
 from nebari_workflow_controller.exceptions import (
@@ -16,6 +19,10 @@
 
 logger = logging.getLogger(__name__)
 
+ARGO_CLIENT_ID = "argo-server-sso"
+# mounted to nebari-workflow-controller deployment as a configmap
+VALID_ARGO_ROLES_CONFIGMAP = "/etc/argo/valid-argo-roles"
+
 
 def process_unhandled_exception(e, return_response, logger):
     error_message = f"An internal error occurred in nebari-workflow-controller while mutating the workflow.  Please open an issue at https://github.com/nebari-dev/nebari-workflow-controller/issues.  The error was: {traceback.format_exc()}"
@@ -45,8 +52,41 @@ def sent_by_argo(workflow: dict):
     return None
 
 
-def get_keycloak_user(request):
-    kcadm = KeycloakAdmin(
+def valid_argo_roles() -> list:
+    with open(VALID_ARGO_ROLES_CONFIGMAP, "r") as f:
+        return json.loads(f.read())
+
+
+def validate_service_account(service_account: str) -> bool:
+    """
+    Check if the service account creating the workflow is from an approved list of service accounts.
+
+    service_account is in the format: "system-serviceaccount-<namespace>-<service account name>"
+    """
+
+    valid_roles = valid_argo_roles()
+    ns = os.environ["NAMESPACE"]
+    sa = service_account.split(f"-{ns}-")
+
+    if len(sa) == 2 and sa[0] == "system-serviceaccount" and sa[1] in valid_roles:
+        return True
+
+    return False
+
+
+def sanitize_label(s: str) -> str:
+    s = s.lower()
+    pattern = r"[^A-Za-z0-9]"
+    return re.sub(pattern, lambda x: "-" + hex(ord(x.group()))[2:], s)
+
+
+def desanitize_label(s: str) -> str:
+    pattern = r"-([A-Za-z0-9][A-Za-z0-9])"
+    return re.sub(pattern, lambda x: chr(int(x.group(1), 16)), s)
+
+
+def create_keycloak_admin() -> KeycloakAdmin:
+    return KeycloakAdmin(
         server_url=os.environ["KEYCLOAK_URL"],
         username=os.environ["KEYCLOAK_USERNAME"],
         password=os.environ["KEYCLOAK_PASSWORD"],
@@ -55,6 +95,10 @@ def get_keycloak_user(request):
         client_id="admin-cli",
     )
 
+
+def get_keycloak_user(request) -> KeycloakUser:
+    kcadm = create_keycloak_admin()
+
     config.incluster_config.load_incluster_config()
 
     keycloak_uid, keycloak_username = get_keycloak_uid_username(
@@ -71,7 +115,7 @@ def get_keycloak_user(request):
 
 
 def get_keycloak_uid_username(
-    kcadm,
+    kcadm: KeycloakAdmin,
     workflow: dict,
     k8s_client: client.ApiClient,
 ) -> KeycloakUser:
@@ -88,8 +132,29 @@ def get_keycloak_uid_username(
 
     if label_added_by_argo == "workflows.argoproj.io/creator":
         keycloak_uid = workflow["metadata"]["labels"][label_added_by_argo]
-        keycloak_username = kcadm.get_user(keycloak_uid)["username"]
-        return keycloak_uid, keycloak_username
+        try:
+            keycloak_username = kcadm.get_user(keycloak_uid)["username"]
+            return keycloak_uid, keycloak_username
+        except KeycloakGetError:
+            logger.warning(
+                f"Keycloak user with UID `{keycloak_uid}` not found. Checking if workflow was created by system-serviceaccount..."
+            )
+            preferred_username = workflow["metadata"]["labels"][
+                "workflows.argoproj.io/creator-preferred-username"
+            ]
+            preferred_username = desanitize_label(preferred_username)
+            if validate_service_account(keycloak_uid):
+                for user in kcadm.get_users():
+                    if user["username"] == preferred_username:
+                        return user["id"], preferred_username
+                raise NWFCUnsupportedException(
+                    "Workflow was created by system-serviceaccount, but user not found in Keycloak. Check that the `PREFERRED_USERNAME` is correctly set in your JupyterLab server."
+                )
+            else:
+                raise NWFCUnsupportedException(
+                    f"Workflow was created by system-serviceaccount submitted by a user without either of the following roles: {valid_argo_roles()}.  Please contact your administrator if you need access."
+                )
+
     elif label_added_by_argo == "workflows.argoproj.io/resubmitted-from-workflow":
         raise NWFCUnsupportedException(
             "Resubmitting workflows is not supported by Nebari Workflow Controller"
@@ -158,13 +223,14 @@ def find_invalid_volume_mount(
                         return denyReason
 
 
-def get_user_pod_spec(keycloak_user):
+def get_user_pod_spec(keycloak_user: KeycloakUser):
     config.incluster_config.load_incluster_config()
     k8s_client = client.CoreV1Api()
 
+    sanitized_username = sanitize_label(keycloak_user.username)
     jupyter_pod_list = k8s_client.list_namespaced_pod(
         os.environ["NAMESPACE"],
-        label_selector=f"hub.jupyter.org/username={keycloak_user.username}",
+        label_selector=f"hub.jupyter.org/username={sanitized_username}",
     ).items
 
     if len(jupyter_pod_list) > 1:
@@ -274,22 +340,30 @@ def mutate_template(
     spec_keep_portions,
     template,
 ):
-    for value, key in container_keep_portions:
-        if "container" not in template:
-            continue
+    target = (
+        "container"
+        if "container" in template
+        else "script"
+        if "script" in template
+        else None
+    )
+
+    if target is None:
+        return
 
+    for value, key in container_keep_portions:
         if isinstance(value, dict):
-            if key in template["container"]:
-                recursive_dict_merge(template["container"][key], value)
+            if key in template[target]:
+                recursive_dict_merge(template[target][key], value)
             else:
-                template["container"][key] = value
+                template[target][key] = value
         elif isinstance(value, list):
-            if key in template["container"]:
-                template["container"][key].append(value)
+            if key in template[target]:
+                template[target][key].extend(value)
             else:
-                template["container"][key] = value
+                template[target][key] = value
         else:
-            template["container"][key] = value
+            template[target][key] = value
 
     for value, key in spec_keep_portions:
         if isinstance(value, dict):
@@ -299,7 +373,7 @@ def mutate_template(
                 template[key] = value
         elif isinstance(value, list):
             if key in template:
-                template[key].append(value)
+                template[key].extend(value)
             else:
                 template[key] = value
         else:

tests/conftest.py

@@ -4,6 +4,7 @@
 
 import pytest
 import yaml
+from keycloak.exceptions import KeycloakGetError
 
 from nebari_workflow_controller.models import KeycloakGroup, KeycloakUser
 
@@ -41,6 +42,28 @@ def all_request_paths(valid_request_paths, invalid_request_paths):
     return valid_request_paths + invalid_request_paths
 
 
+def _valid_service_account_request_paths():
+    return sorted(
+        [
+            str(p)
+            for p in Path("./tests/test_data/requests/valid-service-account").glob(
+                "*.yaml"
+            )
+        ]
+    )
+
+
+def _invalid_service_account_request_paths():
+    return sorted(
+        [
+            str(p)
+            for p in Path("./tests/test_data/requests/invalid-service-account").glob(
+                "*.yaml"
+            )
+        ]
+    )
+
+
 def load_request(request_path):
     with open(request_path) as f:
         return yaml.load(f, Loader=yaml.FullLoader)
@@ -102,3 +125,57 @@ def mocked_get_user_pod_spec(mocker):
         "nebari_workflow_controller.app.get_user_pod_spec",
         return_value=jupyterlab_pod_spec,
     )
+
+
+class MockedKeycloakAdmin:
+    def get_user(self, *args, **kwargs):
+        raise KeycloakGetError("mocked error")
+
+    def get_users(self, *args, **kwargs):
+        return [
+            {
+                "id": "a1234567-abcd-89ef-0123-456789abcdef",
+                "username": "[email protected]",
+                "email": "[email protected]",
+                "firstName": "Valid",
+                "lastName": "User",
+                "enabled": True,
+                "groups": ["admin", "users"],
+            },
+        ]
+
+    def get_user_groups(self, *args, **kwargs):
+        return [
+            {
+                "id": "group123-id",
+                "name": "admin",
+                "path": "/admin",
+                "attributes": {},
+            },
+            {
+                "id": "group456-id",
+                "name": "users",
+                "path": "/users",
+                "attributes": {},
+            },
+        ]
+
+
+VALID_ARGO_ROLES = ["argo-admin", "argo-developer"]
+
+
+@pytest.fixture(scope="function")
+def mock_special_case(mocker):
+    mocker.patch(
+        "nebari_workflow_controller.utils.create_keycloak_admin",
+        return_value=MockedKeycloakAdmin(),
+    )
+    mocker.patch("nebari_workflow_controller.utils.config", return_value=None)
+    mocker.patch(
+        "nebari_workflow_controller.utils.sent_by_argo",
+        return_value="workflows.argoproj.io/creator",
+    )
+    mocker.patch(
+        "nebari_workflow_controller.utils.valid_argo_roles",
+        return_value=VALID_ARGO_ROLES,
+    )

tests/test_app.py

@@ -11,7 +11,13 @@
     get_spec_keep_portions,
     mutate_template,
 )
-from tests.conftest import _invalid_request_paths, _valid_request_paths
+
+from .conftest import (
+    _invalid_request_paths,
+    _invalid_service_account_request_paths,
+    _valid_request_paths,
+    _valid_service_account_request_paths,
+)
 
 
 @pytest.mark.parametrize(
@@ -29,6 +35,24 @@ def test_validate(request_file, allowed, mocked_get_keycloak_user):
         assert response["response"]["status"]["message"]
 
 
+@pytest.mark.parametrize(
+    "request_file,allowed",
+    [(str(p), True) for p in _valid_service_account_request_paths()]
+    + [(str(p), False) for p in _invalid_service_account_request_paths()],
+)
+def test_validate_with_service_account(request_file, allowed, mock_special_case):
+    # The general case is where the user is directly validated against Keycloak.
+    # This is the special case, where the user is validated against Keycloak via a service account.
+
+    with open(request_file) as f:
+        request = yaml.load(f, Loader=yaml.FullLoader)
+    response = validate(request)
+
+    assert response["response"]["allowed"] == allowed
+    if not allowed:
+        assert response["response"]["status"]["message"]
+
+
 def test_mutate_template_doesnt_error(request_templates, jupyterlab_pod_spec):
     api = client.ApiClient()
     container_keep_portions = get_container_keep_portions(jupyterlab_pod_spec, api)