Merge pull request #12 from IsaiahStapleton/assign-class-label

Add assign-class-label mutating webhook

.github/workflows/assign-class-label.yaml

@@ -0,0 +1,43 @@
+name: assign-class-label tests and linting/formatting
+
+on:
+  push:
+    paths:
+    - container-images/assign-class-label/**
+  pull_request:
+    paths:
+    - container-images/assign-class-label/**
+
+jobs:
+  lint-and-test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'  
+
+      - name: Install dependencies
+        working-directory: ./container-images/assign-class-label/
+        run: |
+          pip install -r requirements.txt
+          pip install -r test-requirements.txt
+
+      - name: Run ruff format (formatting)
+        working-directory: ./container-images/assign-class-label/
+        run: |
+          ruff format
+
+      - name: Run ruff check (linting)
+        working-directory: ./container-images/assign-class-label/
+        run: |
+          ruff check --fix
+
+      - name: Run tests
+        working-directory: ./container-images/assign-class-label/
+        run: |
+          pytest tests/

.gitignore

@@ -158,3 +158,8 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+
+webhook.crt 
+webhook.key
+secret.yaml

.ruff.toml

@@ -0,0 +1,18 @@
+[lint]
+# 1. Enable flake8-bugbear (B) rules, in addition to the defaults.
+select = ["E4", "E7", "E9", "F", "B"]
+
+# 2. Avoid enforcing line-length violations (E501)
+ignore = ["E501"]
+
+# 3. Avoid trying to fix flake8-bugbear (B) violations.
+unfixable = ["B"]
+
+# 4. Ignore E402 (import violations) in all __init__.py files, and in selected subdirectories.
+[lint.per-file-ignores]
+"__init__.py" = ["E402"]
+"**/{tests,docs,tools}/*" = ["E402"]
+
+[format]
+# 5. Use single quotes in ruff format.
+quote-style = "single"

README.md

@@ -100,4 +100,71 @@ This script is used to retrieve the URL for a particular notebook associated wit
 	```
 	Enter the notebook name: xxx
 	URL for notebook xxx: xxx
-	```
+	```
+
+## Webhooks
+
+### assign-class-label
+
+This script is used to add labels to the pod of a user denoting which class they belong to (class="classname"). This allows us to differentiate between users of different classes running in the same namespace. This also allows us to create validating [gatekeeper policies](https://github.com/OCP-on-NERC/gatekeeper) for each class.
+
+Before using the assign-class-label webhook, the group-sync cronjob should be run so that the users of the different classes are added to their respective groups in openshift.
+
+In order to modify the deployment follow these steps:
+
+1. Modify the GROUPS env variable to contain the list of classes (openshift groups) of which you would like to assign class labels. This file is found here: webhooks/assign-class-label/deployment.yaml
+
+2. Generate a new OpenSSL certificate
+
+	```
+	openssl req -x509 -sha256 -newkey rsa:2048 -keyout webhook.key -out webhook.crt -days 1024 -nodes -addext "subjectAltName = DNS.1:service_name.namespace.svc"
+	```
+
+	When deployed to rhods-notebooks the command was specified as such:
+
+	```
+	openssl req -x509 -sha256 -newkey rsa:2048 -keyout webhook.key -out webhook.crt -days 1024 -nodes -addext "subjectAltName = DNS.1:assign-class-label-webhook.rhods-notebooks.svc"
+	```
+
+3. Add the cert and key to the required resources:
+
+	```
+	cat webhook.crt | base64 | tr -d '\n'
+	```
+
+	```
+	cat webhook.key | base64 | tr -d '\n'
+	```
+
+	This will encode the certificate and key in base64 format which is required. Copy the output of the webhook.crt to the caBundle in webhooks/assign-class-label/webhook-config.yaml. Then create a secret.yaml that looks like this 
+
+	```
+	apiVersion: v1
+	kind: Secret
+	metadata:
+  		name: webhook-cert
+	type: Opaque
+	data:
+  		webhook.crt: 
+  		webhook.key: 
+	```
+
+	Copy and paste the output of the cat command to the respective fields for webhook.crt and webhook.key. Then execute 
+
+	```
+	oc apply -f secret.yaml --as system:admin
+	```
+
+	within the same namespace that your webhook will be deployed to.
+
+
+4. Change namespace variable in the kubernetes manifests to match namespace you want the webhook to be deployed to.
+
+5. From webhooks/assign-class-label/ directory run:
+```
+	oc apply -k . --as system:admin
+```
+
+***Steps 2, 3, and 4 are only required if you are deploying to a new namespace/environment.***
+
+The python script and docker image used for the webserver should not need changes made to it. But in the case that changes must be made, the Dockerfile and python script can be found at docker/src/python/assign-class-label/.

container-images/assign-class-label/Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.12-slim
+
+WORKDIR /app/
+
+COPY requirements.txt ./
+
+RUN pip install -r requirements.txt
+
+COPY . ./
+
+EXPOSE 5000
+
+CMD ["gunicorn", "wsgi:webhook", "--log-level=info", "--workers", "3", "--bind", "0.0.0.0:5000", "--keyfile", "/certs/webhook.key", "--certfile", "/certs/webhook.crt"]

container-images/assign-class-label/models.py

@@ -0,0 +1,37 @@
+from pydantic import BaseModel, constr
+from typing import Dict, Optional
+
+
+class PodMetadata(BaseModel):
+    labels: Optional[Dict[str, str]] = None
+
+
+class PodObject(BaseModel):
+    metadata: PodMetadata
+
+
+class AdmissionRequest(BaseModel):
+    uid: constr(min_length=1)
+    object: PodObject
+
+
+class AdmissionReview(BaseModel):
+    request: AdmissionRequest
+
+
+class Status(BaseModel):
+    message: Optional[str] = None
+
+
+class AdmissionResponse(BaseModel):
+    uid: str
+    allowed: bool
+    status: Optional[Status] = None
+    patchType: Optional[str] = None
+    patch: Optional[str] = None
+
+
+class AdmissionReviewResponse(BaseModel):
+    apiVersion: str = 'admission.k8s.io/v1'
+    kind: str = 'AdmissionReview'
+    response: AdmissionResponse

container-images/assign-class-label/mutate.py

@@ -0,0 +1,165 @@
+import logging
+import json
+import base64
+from flask import Flask, request, jsonify, Response
+from kubernetes import config, client
+from openshift.dynamic import DynamicClient
+
+from pydantic import ValidationError
+from typing import Any, List
+
+from models import AdmissionReviewResponse, AdmissionReview, AdmissionResponse, Status
+
+LOG = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def decode_pod_user(pod_user: str) -> str:
+    return pod_user.replace('-40', '@').replace('-2e', '.')
+
+
+def get_client() -> DynamicClient:
+    try:
+        config.load_config()
+        k8s_client = client.ApiClient()
+        dyn_client = DynamicClient(k8s_client)
+        return dyn_client
+    except config.ConfigException as e:
+        LOG.error('Could not configure Kubernetes client: %s', str(e))
+        exit(1)
+
+
+def get_group_resource(dyn_client):
+    return dyn_client.resources.get(api_version='user.openshift.io/v1', kind='Group')
+
+
+# Get users of a given group
+def get_group_members(group_resource: Any, group_name: str) -> List[str]:
+    group_obj = group_resource.get(name=group_name)
+    return group_obj.users
+
+
+def assign_class_label(
+    pod: dict[str, Any], groups: list[str], dyn_client: DynamicClient
+) -> str | None:
+    # Extract pod metadata
+    try:
+        pod_metadata = pod.get('metadata', {})
+        pod_labels = pod_metadata.get('labels', {})
+        pod_user = pod_labels.get('opendatahub.io/user', None)
+    except AttributeError as e:
+        LOG.error(f'Error extracting pod information: {e}')
+        return None
+
+    if pod_user is None:
+        return None
+
+    pod_user = decode_pod_user(pod_user)
+
+    group_resource = get_group_resource(dyn_client)
+
+    # Iterate through classes
+    for group in groups:
+        users = get_group_members(group_resource, group)
+
+        # Check if group has no users
+        if not users:
+            LOG.warning(f'Group {group} has no users or users attribute is not a list.')
+            continue
+
+        # Compare users in the groups (classes) with the pod user
+        if pod_user in users:
+            LOG.info(f'Assigning class label: {group} to user {pod_user}')
+            return group
+
+    return None
+
+
+def create_app(**config: Any) -> Flask:
+    app = Flask(__name__)
+    app.config.from_prefixed_env('RHOAI_CLASS')
+    app.config.update(config)
+
+    if not app.config['GROUPS']:
+        LOG.error('RHOAI_CLASS_GROUPS environment variables are required.')
+        exit(1)
+
+    groups = app.config['GROUPS'].split(',')
+
+    dyn_client = get_client()
+
+    @app.route('/mutate', methods=['POST'])
+    def mutate_pod() -> Response:
+        # Grab pod for mutation and validate request
+        try:
+            admission_review = AdmissionReview(**request.get_json())
+        except ValidationError as e:
+            LOG.error('Validation error: %s', e)
+            return (
+                jsonify(
+                    AdmissionReviewResponse(
+                        response=AdmissionResponse(
+                            uid=request.json.get('request', {}).get('uid', ''),
+                            allowed=False,
+                            status=Status(message=f'Invalid request: {e}'),
+                        )
+                    ).model_dump()
+                ),
+                400,
+                {'content-type': 'application/json'},
+            )
+
+        uid = admission_review.request.uid
+        pod = admission_review.request.object.model_dump()
+
+        # Grab class that the pod user belongs to
+        try:
+            class_label = assign_class_label(pod, groups, dyn_client)
+        except Exception as err:
+            LOG.error('failed to assign class label: %s', err)
+            return 'unexpected error encountered', 500, {'content-type': 'text/plain'}
+
+        # If user not in any class, return without modifications
+        if not class_label:
+            return (
+                jsonify(
+                    AdmissionReviewResponse(
+                        response=AdmissionResponse(
+                            uid=uid,
+                            allowed=True,
+                            status=Status(message='No class label assigned.'),
+                        )
+                    ).model_dump()
+                ),
+                200,
+                {'content-type': 'application/json'},
+            )
+
+        # Generate JSON Patch to add class label
+        patch = [
+            {
+                'op': 'add',
+                'path': '/metadata/labels/nerc.mghpcc.org~1class',
+                'value': class_label,
+            }
+        ]
+
+        # Encode patch as base64 for response
+        patch_base64 = base64.b64encode(json.dumps(patch).encode('utf-8')).decode(
+            'utf-8'
+        )
+
+        # Return webhook response that includes the patch to add class label
+        return (
+            jsonify(
+                AdmissionReviewResponse(
+                    response=AdmissionResponse(
+                        uid=uid, allowed=True, patchType='JSONPatch', patch=patch_base64
+                    )
+                ).model_dump()
+            ),
+            200,
+            {'content-type': 'application/json'},
+        )
+
+    return app

container-images/assign-class-label/requirements.txt

@@ -0,0 +1,5 @@
+flask
+kubernetes
+openshift
+gunicorn
+pydantic

container-images/assign-class-label/test-requirements.txt

@@ -0,0 +1,2 @@
+pytest
+ruff