paasta-tools 1.21.3__py3-none-any.whl

paasta_tools/kubernetes/remote_run.py

@@ -0,0 +1,558 @@
+# Copyright 2015-2016 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import hashlib
+import logging
+from time import sleep
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import TypedDict
+
+from kubernetes.client import AuthenticationV1TokenRequest
+from kubernetes.client import V1Job
+from kubernetes.client import V1ObjectMeta
+from kubernetes.client import V1Pod
+from kubernetes.client import V1PolicyRule
+from kubernetes.client import V1Role
+from kubernetes.client import V1RoleBinding
+from kubernetes.client import V1RoleRef
+from kubernetes.client import V1ServiceAccount
+from kubernetes.client import V1Subject
+from kubernetes.client import V1TokenRequestSpec
+from kubernetes.client.exceptions import ApiException
+
+from paasta_tools.adhoc_tools import load_adhoc_job_config
+from paasta_tools.eks_tools import EksDeploymentConfig
+from paasta_tools.eks_tools import load_eks_service_config
+from paasta_tools.kubernetes.application.controller_wrappers import (
+    get_application_wrapper,
+)
+from paasta_tools.kubernetes_tools import get_all_service_accounts
+from paasta_tools.kubernetes_tools import JOB_TYPE_LABEL_NAME
+from paasta_tools.kubernetes_tools import KubeClient
+from paasta_tools.kubernetes_tools import limit_size_with_hash
+from paasta_tools.kubernetes_tools import paasta_prefixed
+from paasta_tools.utils import load_system_paasta_config
+
+
+logger = logging.getLogger(__name__)
+REMOTE_RUN_JOB_LABEL = "remote-run"
+POD_OWNER_LABEL = paasta_prefixed("pod_owner")
+TOOLBOX_MOCK_SERVICE = "prod-toolbox"
+DEFAULT_MAX_DURATION_LIMIT = 8 * 60 * 60  # 8 hours
+
+
+class RemoteRunError(Exception):
+    pass
+
+
+class RemoteRunOutcome(TypedDict, total=False):
+    status: int
+    message: str
+    job_name: str
+    pod_name: str
+    pod_address: str
+    namespace: str
+
+
+def _format_remote_run_job_name(
+    job: V1Job,
+    user: str,
+) -> str:
+    """Format name for remote run job
+
+    :param V1Job job: job definition
+    :param str user: the user requesting the remote-run
+    :return: job name
+    """
+    return limit_size_with_hash(f"remote-run-{user}-{job.metadata.name}")
+
+
+def remote_run_start(
+    service: str,
+    instance: str,
+    cluster: str,
+    user: str,
+    interactive: bool,
+    recreate: bool,
+    max_duration: int,
+    is_toolbox: bool,
+) -> RemoteRunOutcome:
+    """Trigger remote-run job
+
+    :param str service: service name
+    :param str instance: service instance
+    :param str cluster: paasta cluster
+    :param str user: the user requesting the remote-run sandbox
+    :param bool interactive: whether it is expected to access the remote-run job interactively
+    :param bool recreate: whether to recreate remote-run job if existing
+    :param int max_duration: maximum allowed duration for the remote-ruh job
+    :param bool is_toolbox: requested job is for a toolbox container
+    :return: outcome of the operation, and resulting Kubernetes pod information
+    """
+    kube_client = KubeClient()
+
+    # Load the service deployment settings
+    deployment_config = (
+        generate_toolbox_deployment(service, cluster, user)
+        if is_toolbox
+        else load_eks_service_config(service, instance, cluster)
+    )
+
+    # Set to interactive mode
+    if interactive and not is_toolbox:
+        deployment_config.config_dict["cmd"] = f"sleep {max_duration}"
+
+    # Create the app with a new name
+    formatted_job = deployment_config.format_kubernetes_job(
+        job_label=REMOTE_RUN_JOB_LABEL,
+        deadline_seconds=max_duration,
+        keep_routable_ip=is_toolbox,
+    )
+    job_name = _format_remote_run_job_name(formatted_job, user)
+    formatted_job.metadata.name = job_name
+    app_wrapper = get_application_wrapper(formatted_job)
+    app_wrapper.soa_config = deployment_config
+
+    # Launch pod
+    logger.info(f"Starting {job_name}")
+    try:
+        app_wrapper.create(kube_client)
+    except ApiException as e:
+        if e.status != 409:
+            raise
+        if recreate:
+            remote_run_stop(
+                service=service,
+                instance=instance,
+                cluster=cluster,
+                user=user,
+                is_toolbox=is_toolbox,
+            )
+            return remote_run_start(
+                service=service,
+                instance=instance,
+                cluster=cluster,
+                user=user,
+                interactive=interactive,
+                recreate=False,
+                max_duration=max_duration,
+                is_toolbox=is_toolbox,
+            )
+
+    return {
+        "status": 200,
+        "message": "Remote run sandbox started",
+        "job_name": job_name,
+    }
+
+
+def remote_run_ready(
+    service: str,
+    instance: str,
+    cluster: str,
+    job_name: str,
+    user: str,
+    is_toolbox: bool,
+) -> RemoteRunOutcome:
+    """Check if remote-run pod is ready
+
+    :param str service: service name
+    :param str instance: service instance
+    :param str cluster: paasta cluster
+    :param str job_name: name of the remote-run job to check
+    :param bool is_toolbox: requested job is for a toolbox container
+    :return: job status, with pod info
+    """
+    kube_client = KubeClient()
+
+    # Load the service deployment settings
+    deployment_config = (
+        generate_toolbox_deployment(service, cluster, user)
+        if is_toolbox
+        else load_eks_service_config(service, instance, cluster)
+    )
+    namespace = deployment_config.get_namespace()
+
+    pod = find_job_pod(kube_client, namespace, job_name)
+    if not pod:
+        return {"status": 404, "message": "No pod found"}
+    if pod.status.phase == "Running":
+        result: RemoteRunOutcome = {
+            "status": 200,
+            "message": "Pod ready",
+            "pod_name": pod.metadata.name,
+            "namespace": namespace,
+        }
+        if is_toolbox:
+            result["pod_address"] = pod.status.pod_ip
+        return result
+    return {
+        "status": 204,
+        "message": "Pod not ready",
+    }
+
+
+def remote_run_stop(
+    service: str,
+    instance: str,
+    cluster: str,
+    user: str,
+    is_toolbox: bool,
+) -> RemoteRunOutcome:
+    """Stop remote-run job
+
+    :param str service: service name
+    :param str instance: service instance
+    :param str cluster: paasta cluster
+    :param str user: the user requesting the remote-run sandbox
+    :param bool is_toolbox: requested job is for a toolbox container
+    :return: outcome of the operation
+    """
+    kube_client = KubeClient()
+
+    # Load the service deployment settings
+    deployment_config = (
+        generate_toolbox_deployment(service, cluster, user)
+        if is_toolbox
+        else load_eks_service_config(service, instance, cluster)
+    )
+
+    # Rebuild the job metadata
+    formatted_job = deployment_config.format_kubernetes_job(
+        job_label=REMOTE_RUN_JOB_LABEL
+    )
+    job_name = _format_remote_run_job_name(formatted_job, user)
+    formatted_job.metadata.name = job_name
+
+    # Stop the job
+    logger.info(f"Stopping {job_name}")
+    app_wrapper = get_application_wrapper(formatted_job)
+    app_wrapper.soa_config = deployment_config
+    app_wrapper.deep_delete(kube_client)
+
+    return {"status": 200, "message": "Job successfully removed"}
+
+
+def remote_run_token(
+    service: str,
+    instance: str,
+    cluster: str,
+    user: str,
+) -> str:
+    """Creates a short lived token for execing into a pod
+
+    :param str service: service name
+    :param str instance: service instance
+    :param str cluster: paasta cluster
+    :param str user: the user requesting the remote-run sandbox
+    """
+    kube_client = KubeClient()
+
+    # Load the service deployment settings
+    deployment_config = load_eks_service_config(service, instance, cluster)
+    namespace = deployment_config.get_namespace()
+
+    # Rebuild the job metadata
+    formatted_job = deployment_config.format_kubernetes_job(
+        job_label=REMOTE_RUN_JOB_LABEL
+    )
+    job_name = _format_remote_run_job_name(formatted_job, user)
+
+    # Find pod and create exec token for it
+    pod = find_job_pod(kube_client, namespace, job_name)
+    if not pod:
+        raise RemoteRunError(f"Pod for {job_name} not found")
+    pod_name = pod.metadata.name
+    logger.info(f"Generating temporary service account token for {pod_name}")
+    service_account = create_remote_run_service_account(
+        kube_client, namespace, pod_name, user
+    )
+    role = create_pod_scoped_role(kube_client, namespace, pod_name, user)
+    bind_role_to_service_account(kube_client, namespace, service_account, role, user)
+    return create_temp_exec_token(kube_client, namespace, service_account)
+
+
+def generate_toolbox_deployment(
+    service: str, cluster: str, user: str
+) -> EksDeploymentConfig:
+    """Creates virtual EKS deployment for toolbox containers starting from adhoc configuration
+
+    :param str service: toolbox name
+    :param str cluster: target deployment cluster
+    :param str user: user requesting the toolbox
+    :return: deployment configuration
+    """
+    if not user.isalnum():
+        raise RemoteRunError(
+            f"Provided username contains non-alphanumeric characters: {user}"
+        )
+    # NOTE: API authorization is enforced by service, and we want different rules
+    # for each toolbox, so clients send a combined service-instance string, and then
+    # we split it here to load the correct instance settings.
+    adhoc_instance = service[len(TOOLBOX_MOCK_SERVICE) + 1 :]
+    adhoc_deployment = load_adhoc_job_config(
+        TOOLBOX_MOCK_SERVICE,
+        adhoc_instance,
+        cluster,
+        load_deployments=False,
+    )
+    # NOTE: we're explicitly dynamically mounting a single user's public keys
+    # as we want these pods to only be usable by said user.
+    adhoc_deployment.config_dict.setdefault("extra_volumes", []).append(
+        {
+            "containerPath": f"/etc/authorized_keys.d/{user}.pub",
+            "hostPath": f"/etc/authorized_keys.d/{user}.pub",
+            "mode": "RO",
+        },
+    )
+    adhoc_deployment.config_dict.setdefault("env", {})["SANDBOX_USER"] = user
+    adhoc_deployment.config_dict["routable_ip"] = True
+    return EksDeploymentConfig(
+        service=service,
+        cluster=cluster,
+        instance="main",
+        config_dict=adhoc_deployment.config_dict,
+        branch_dict=adhoc_deployment.branch_dict,
+    )
+
+
+def find_job_pod(
+    kube_client: KubeClient,
+    namespace: str,
+    job_name: str,
+    job_label: str = REMOTE_RUN_JOB_LABEL,
+    retries: int = 3,
+) -> Optional[V1Pod]:
+    """Locate pod for remote-run job
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: the pod namespace
+    :param str job_name: remote-run job name
+    :param int retries: maximum number of attemps
+    :return: pod object if found
+    """
+    selectors = (
+        f"{paasta_prefixed(JOB_TYPE_LABEL_NAME)}={job_label}",
+        f"job-name={job_name}",
+    )
+    for _ in range(retries):
+        pod_list = kube_client.core.list_namespaced_pod(
+            namespace,
+            label_selector=",".join(selectors),
+        )
+        if pod_list.items:
+            return pod_list.items[0]
+        sleep(0.5)
+    return None
+
+
+def create_temp_exec_token(
+    kube_client: KubeClient,
+    namespace: str,
+    service_account: str,
+) -> str:
+    """Create a short lived token for service account
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: service account namespace
+    :param str service_account: service account name
+    :return: token value
+    """
+    token_spec = V1TokenRequestSpec(
+        expiration_seconds=600,  # minimum allowed by k8s
+        audiences=[],
+    )
+    request = AuthenticationV1TokenRequest(spec=token_spec)
+    response = kube_client.core.create_namespaced_service_account_token(
+        service_account, namespace, request
+    )
+    return response.status.token
+
+
+def get_remote_run_service_accounts(
+    kube_client: KubeClient, namespace: str, user: str = ""
+) -> Sequence[V1ServiceAccount]:
+    """List all temporary service account related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :param str user: optionally filter by owning user
+    :return: list of service accounts
+    """
+    return get_all_service_accounts(
+        kube_client,
+        namespace=namespace,
+        label_selector=(f"{POD_OWNER_LABEL}={user}" if user else POD_OWNER_LABEL),
+    )
+
+
+def create_remote_run_service_account(
+    kube_client: KubeClient,
+    namespace: str,
+    pod_name: str,
+    user: str,
+) -> str:
+    """Create service account to exec into remote-run pod
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :param str pod_name: pod name
+    :param str user: user requiring credentials
+    """
+    pod_name_hash = hashlib.sha1(pod_name.encode("utf-8")).hexdigest()[:12]
+    service_account_name = limit_size_with_hash(f"remote-run-{user}-{pod_name_hash}")
+    service_accounts = get_remote_run_service_accounts(kube_client, namespace, user)
+    if any(item.metadata.name == service_account_name for item in service_accounts):
+        return service_account_name
+    service_account = V1ServiceAccount(
+        metadata=V1ObjectMeta(
+            name=service_account_name,
+            namespace=namespace,
+            labels={POD_OWNER_LABEL: user},
+        )
+    )
+    kube_client.core.create_namespaced_service_account(
+        namespace=namespace, body=service_account
+    )
+    return service_account_name
+
+
+def create_pod_scoped_role(
+    kube_client: KubeClient,
+    namespace: str,
+    pod_name: str,
+    user: str,
+) -> str:
+    """Create role with execution access to specific pod
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :param str pod_name: pod name
+    :param str user: user requiring the role
+    :return: name of the role
+    """
+    pod_name_hash = hashlib.sha1(pod_name.encode("utf-8")).hexdigest()[:12]
+    role_name = f"remote-run-role-{pod_name_hash}"
+    policy = V1PolicyRule(
+        verbs=["create", "get"],
+        resources=["pods", "pods/exec"],
+        resource_names=[pod_name],
+        api_groups=[""],
+    )
+    role = V1Role(
+        rules=[policy],
+        metadata=V1ObjectMeta(
+            name=role_name,
+            labels={POD_OWNER_LABEL: user},
+        ),
+    )
+    try:
+        kube_client.rbac.create_namespaced_role(namespace=namespace, body=role)
+    except ApiException as e:
+        if e.status != 409:
+            raise
+    return role_name
+
+
+def bind_role_to_service_account(
+    kube_client: KubeClient,
+    namespace: str,
+    service_account: str,
+    role: str,
+    user: str,
+) -> None:
+    """Bind service account to role
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: service account namespace
+    :param str service_account: service account name
+    :param str role: role name
+    :param str user: user requiring the role
+    """
+    role_binding = V1RoleBinding(
+        metadata=V1ObjectMeta(
+            name=limit_size_with_hash(f"remote-run-binding-{role}"),
+            namespace=namespace,
+            labels={POD_OWNER_LABEL: user},
+        ),
+        role_ref=V1RoleRef(
+            api_group="rbac.authorization.k8s.io",
+            kind="Role",
+            name=role,
+        ),
+        subjects=[
+            V1Subject(
+                kind="ServiceAccount",
+                name=service_account,
+            ),
+        ],
+    )
+    try:
+        kube_client.rbac.create_namespaced_role_binding(
+            namespace=namespace,
+            body=role_binding,
+        )
+    except ApiException as e:
+        if e.status != 409:
+            raise
+
+
+def get_remote_run_roles(kube_client: KubeClient, namespace: str) -> List[V1Role]:
+    """List all temporary roles related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: role namespace
+    :return: list of roles
+    """
+    return kube_client.rbac.list_namespaced_role(
+        namespace,
+        label_selector=POD_OWNER_LABEL,
+    ).items
+
+
+def get_remote_run_role_bindings(
+    kube_client: KubeClient, namespace: str
+) -> List[V1RoleBinding]:
+    """List all temporary role bindings related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: role namespace
+    :return: list of roles
+    """
+    return kube_client.rbac.list_namespaced_role_binding(
+        namespace,
+        label_selector=POD_OWNER_LABEL,
+    ).items
+
+
+def get_remote_run_jobs(kube_client: KubeClient, namespace: str) -> List[V1Job]:
+    """List all remote-run jobs
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: job namespace
+    """
+    return kube_client.batches.list_namespaced_job(
+        namespace,
+        label_selector=f"{paasta_prefixed(JOB_TYPE_LABEL_NAME)}={REMOTE_RUN_JOB_LABEL}",
+    ).items
+
+
+def get_max_job_duration_limit() -> int:
+    """Get maximum configured duration for a remote run job
+
+    :return: max duration in seconds
+    """
+    system_config = load_system_paasta_config()
+    return system_config.get_remote_run_duration_limit(DEFAULT_MAX_DURATION_LIMIT)