DataOps-platform/app/runner/auth.py

"""Short-lived signed task tokens bound to one immutable workflow node."""

from __future__ import annotations

import hashlib
import base64
import json
import time
import uuid
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any

import jwt

ISSUER = "dataops-platform"
AUDIENCE = "dataops-runner"
ALGORITHM = "HS256"


class TaskTokenInvalid(ValueError):
    pass


class TaskTokenExpired(TaskTokenInvalid):
    pass


def node_digest(node: Mapping[str, Any]) -> str:
    canonical = json.dumps(
        node,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _secret_bytes(secret: str) -> bytes:
    encoded = str(secret or "").encode("utf-8")
    if len(encoded) < 32:
        raise ValueError("task token secret must be at least 32 bytes")
    return encoded


@dataclass(frozen=True)
class TaskClaims:
    task_uid: str
    dataflow_uid: str
    deployment_id: str
    environment: str
    workflow_version: int
    correlation_id: str
    node_id: str
    node_type: str
    node_digest: str
    purpose: str
    write_authorized: bool
    issued_at: int
    expires_at: int
    jti: str


class TaskTokenIssuer:
    def __init__(self, secret, *, clock=None, ttl_seconds=60):
        self._secret = _secret_bytes(secret)
        self._clock = clock or time.time
        self._ttl_seconds = int(ttl_seconds)
        if self._ttl_seconds < 1 or self._ttl_seconds > 300:
            raise ValueError("task token ttl must be between 1 and 300 seconds")

    def issue(
        self,
        *,
        task_uid,
        dataflow_uid,
        deployment_id,
        environment,
        workflow_version,
        correlation_id,
        node,
        write_authorized=False,
    ):
        now = int(self._clock())
        if environment not in {"development", "test", "production"}:
            raise ValueError("task environment is invalid")
        purpose = str(node.get("purpose") or "read")
        if purpose == "write" and not write_authorized:
            raise ValueError("write task requires trusted write authorization")
        payload = {
            "iss": ISSUER,
            "aud": AUDIENCE,
            "sub": str(task_uid),
            "task_uid": str(task_uid),
            "dataflow_uid": str(dataflow_uid),
            "deployment_id": str(deployment_id),
            "environment": str(environment),
            "workflow_version": int(workflow_version),
            "correlation_id": str(correlation_id),
            "node_id": str(node.get("id") or ""),
            "node_type": str(node.get("type") or ""),
            "node_digest": node_digest(node),
            "purpose": purpose,
            "write_authorized": bool(write_authorized),
            "iat": now,
            "exp": now + self._ttl_seconds,
            "jti": str(uuid.uuid4()),
        }
        return jwt.encode(payload, self._secret, algorithm=ALGORITHM)


class TaskTokenVerifier:
    def __init__(
        self,
        secret,
        *,
        clock=None,
        replay_window_seconds=900,
    ):
        self._secret = _secret_bytes(secret)
        self._clock = clock or time.time
        self._replay_window_seconds = int(replay_window_seconds)
        if (
            self._replay_window_seconds < 1
            or self._replay_window_seconds > 3600
        ):
            raise ValueError("task replay window is invalid")

    def _verify(self, token, *, node, replay_only):
        try:
            segments = str(token).split(".")
            if len(segments) != 3:
                raise ValueError("segment count")
            for segment in segments:
                decoded = base64.urlsafe_b64decode(
                    segment + "=" * (-len(segment) % 4)
                )
                canonical = base64.urlsafe_b64encode(decoded).rstrip(b"=").decode("ascii")
                if canonical != segment:
                    raise ValueError("non-canonical base64url")
        except (TypeError, ValueError, UnicodeError) as exc:
            raise TaskTokenInvalid("task token is invalid") from exc
        try:
            payload = jwt.decode(
                token,
                self._secret,
                algorithms=[ALGORITHM],
                audience=AUDIENCE,
                issuer=ISSUER,
                options={"verify_exp": False, "verify_iat": False},
            )
        except jwt.PyJWTError as exc:
            raise TaskTokenInvalid("task token is invalid") from exc

        required = {
            "task_uid",
            "dataflow_uid",
            "deployment_id",
            "environment",
            "workflow_version",
            "correlation_id",
            "node_id",
            "node_type",
            "node_digest",
            "purpose",
            "iat",
            "exp",
            "jti",
        }
        if required - set(payload):
            raise TaskTokenInvalid("task token claims are incomplete")
        now = int(self._clock())
        try:
            issued_at = int(payload["iat"])
            expires_at = int(payload["exp"])
            workflow_version = int(payload["workflow_version"])
        except (TypeError, ValueError) as exc:
            raise TaskTokenInvalid("task token claims are invalid") from exc
        if issued_at > now + 5 or expires_at - issued_at > 300:
            raise TaskTokenInvalid("task token lifetime is invalid")
        if expires_at < now:
            if not replay_only:
                raise TaskTokenExpired("task token has expired")
            if now - expires_at > self._replay_window_seconds:
                raise TaskTokenExpired(
                    "task token replay window has expired"
                )
        if node_digest(node) != payload["node_digest"]:
            raise TaskTokenInvalid("task token node binding does not match")
        if (
            payload["node_id"] != node.get("id")
            or payload["node_type"] != node.get("type")
            or payload["purpose"] != str(node.get("purpose") or "read")
        ):
            raise TaskTokenInvalid("task token node binding does not match")
        if payload["environment"] not in {
            "development",
            "test",
            "production",
        }:
            raise TaskTokenInvalid("task token environment is invalid")
        return TaskClaims(
            task_uid=str(payload["task_uid"]),
            dataflow_uid=str(payload["dataflow_uid"]),
            deployment_id=str(payload["deployment_id"]),
            environment=str(payload["environment"]),
            workflow_version=workflow_version,
            correlation_id=str(payload["correlation_id"]),
            node_id=str(payload["node_id"]),
            node_type=str(payload["node_type"]),
            node_digest=str(payload["node_digest"]),
            purpose=str(payload["purpose"]),
            write_authorized=bool(payload.get("write_authorized", False)),
            issued_at=issued_at,
            expires_at=expires_at,
            jti=str(payload["jti"]),
        )

    def verify(self, token, *, node):
        return self._verify(token, node=node, replay_only=False)

    def verify_for_replay(self, token, *, node):
        return self._verify(token, node=node, replay_only=True)