from __future__ import annotations from datetime import UTC, datetime from typing import Any from kubernetes import client, config from kubernetes.client.exceptions import ApiException from .config import Settings from .errors import PanelError from .models import AccountBase, InstanceSummary MANAGED_LABEL = "app.kubernetes.io/name=userbot" INSTANCE_LABEL = "app.kubernetes.io/instance" MANAGED_BY_LABEL = "app.kubernetes.io/managed-by" DISPLAY_ANNOTATION = "userbot.forust.xyz/display-name" LEGACY_ANNOTATION = "userbot.forust.xyz/legacy" CREDENTIALS_ANNOTATION = "userbot.forust.xyz/credentials-secret" PVC_ANNOTATION = "userbot.forust.xyz/pvc" RESTART_ANNOTATION = "userbot.forust.xyz/restarted-at" def _selector(labels: dict[str, str] | None) -> str: return ",".join(f"{key}={value}" for key, value in (labels or {}).items()) def _as_datetime(value: Any) -> datetime | None: if value is None: return None if isinstance(value, datetime): return value return getattr(value, "replace", lambda **_: None)(tzinfo=UTC) class KubernetesService: def __init__( self, settings: Settings, *, core: client.CoreV1Api | None = None, apps: client.AppsV1Api | None = None, custom: client.CustomObjectsApi | None = None, ) -> None: self.settings = settings if core is None or apps is None: try: config.load_incluster_config() except config.ConfigException: config.load_kube_config() self.core = core or client.CoreV1Api() self.apps = apps or client.AppsV1Api() self.custom = custom or client.CustomObjectsApi() def ensure_prerequisites(self) -> None: try: self.core.read_namespaced_secret( self.settings.common_secret, self.settings.namespace, ) self.core.read_namespaced_config_map( self.settings.common_config, self.settings.namespace, ) except ApiException as exc: if exc.status == 404: raise PanelError( 503, "Userbot common Secret or ConfigMap is missing in the userbot namespace", ) from exc raise self._api_error(exc, "Could not verify userbot prerequisites") from exc def list_instances( self, query: str = "", status: str = "", ) -> list[InstanceSummary]: instances: list[InstanceSummary] = [] for namespace in (self.settings.namespace, *self.settings.legacy_namespaces): try: deployments = self.apps.list_namespaced_deployment( namespace, label_selector=MANAGED_LABEL, ).items except ApiException as exc: raise self._api_error(exc, f"Could not list Deployments in {namespace}") from exc instances.extend(self._summarize(namespace, deployment) for deployment in deployments) query = query.strip().lower() if query: instances = [ item for item in instances if query in item.instance_id.lower() or query in item.display_name.lower() or query in (item.pod or "").lower() ] if status: instances = [item for item in instances if item.status == status] return sorted(instances, key=lambda item: (item.legacy, item.instance_id)) def get_instance(self, instance_id: str) -> InstanceSummary: namespace, deployment = self._find_deployment(instance_id) return self._summarize(namespace, deployment) def assert_available(self, instance_id: str) -> None: names = self._resource_names(instance_id) checks = ( (self.apps.read_namespaced_deployment, names["deployment"], "Deployment"), (self.core.read_namespaced_secret, names["secret"], "Secret"), (self.core.read_namespaced_persistent_volume_claim, names["pvc"], "PVC"), ) for read, name, kind in checks: try: read(name, self.settings.namespace) except ApiException as exc: if exc.status == 404: continue raise self._api_error(exc, f"Could not check {kind} {name}") from exc raise PanelError(409, f"{kind} {name} already exists") def provision(self, account: AccountBase, session_string: str) -> InstanceSummary: self.ensure_prerequisites() self.assert_available(account.instance_id) names = self._resource_names(account.instance_id) created: list[tuple[str, str]] = [] try: self.core.create_namespaced_secret( self.settings.namespace, self._secret(account, session_string, names), ) created.append(("secret", names["secret"])) self.core.create_namespaced_persistent_volume_claim( self.settings.namespace, self._pvc(account, names), ) created.append(("pvc", names["pvc"])) self.apps.create_namespaced_deployment( self.settings.namespace, self._deployment(account, names), ) created.append(("deployment", names["deployment"])) except ApiException as exc: self._rollback(created) raise self._api_error(exc, "Could not create userbot instance") from exc return self.get_instance(account.instance_id) def scale(self, instance_id: str, replicas: int) -> InstanceSummary: namespace, deployment = self._find_deployment(instance_id) try: self.apps.patch_namespaced_deployment_scale( deployment.metadata.name, namespace, {"spec": {"replicas": replicas}}, ) except ApiException as exc: raise self._api_error(exc, "Could not scale userbot instance") from exc return self.get_instance(instance_id) def restart(self, instance_id: str) -> InstanceSummary: namespace, deployment = self._find_deployment(instance_id) if (deployment.spec.replicas or 0) == 0: raise PanelError(409, "Stopped instance cannot be restarted") timestamp = datetime.now(UTC).isoformat() try: self.apps.patch_namespaced_deployment( deployment.metadata.name, namespace, { "spec": { "template": { "metadata": { "annotations": {RESTART_ANNOTATION: timestamp}, } } } }, ) except ApiException as exc: raise self._api_error(exc, "Could not restart userbot instance") from exc return self.get_instance(instance_id) def logs(self, instance_id: str, tail: int = 250) -> str: namespace, deployment = self._find_deployment(instance_id) pods = self._pods_for_deployment(namespace, deployment) if not pods: raise PanelError(409, "Userbot Pod is not running") pod = pods[0] container = deployment.spec.template.spec.containers[0].name try: return self.core.read_namespaced_pod_log( pod.metadata.name, namespace, container=container, tail_lines=max(1, min(tail, 1000)), timestamps=True, ) except ApiException as exc: raise self._api_error(exc, "Could not read userbot logs") from exc def delete(self, instance_id: str, *, delete_data: bool) -> None: namespace, deployment = self._find_deployment(instance_id) annotations = deployment.metadata.annotations or {} if annotations.get(LEGACY_ANNOTATION) == "true" or namespace != self.settings.namespace: raise PanelError(409, "Legacy instances cannot be deleted from the panel") names = self._resource_names(instance_id) secret_name = annotations.get(CREDENTIALS_ANNOTATION, names["secret"]) pvc_name = annotations.get(PVC_ANNOTATION, names["pvc"]) operations = [ ( self.apps.delete_namespaced_deployment, (deployment.metadata.name, namespace), {"propagation_policy": "Foreground"}, ), (self.core.delete_namespaced_secret, (secret_name, namespace), {}), ] if delete_data: operations.append( ( self.core.delete_namespaced_persistent_volume_claim, (pvc_name, namespace), {}, ) ) for delete_resource, args, kwargs in operations: try: delete_resource(*args, **kwargs) except ApiException as exc: if exc.status != 404: raise self._api_error(exc, "Could not delete userbot instance") from exc def _find_deployment(self, instance_id: str) -> tuple[str, Any]: selector = f"{MANAGED_LABEL},{INSTANCE_LABEL}={instance_id}" for namespace in (self.settings.namespace, *self.settings.legacy_namespaces): try: items = self.apps.list_namespaced_deployment( namespace, label_selector=selector, ).items except ApiException as exc: raise self._api_error(exc, "Could not find userbot instance") from exc if items: return namespace, items[0] raise PanelError(404, f"Instance {instance_id} does not exist") def _summarize(self, namespace: str, deployment: Any) -> InstanceSummary: labels = deployment.metadata.labels or {} annotations = deployment.metadata.annotations or {} instance_id = labels.get(INSTANCE_LABEL, deployment.metadata.name) legacy = annotations.get(LEGACY_ANNOTATION) == "true" pods = self._pods_for_deployment(namespace, deployment) pod = pods[0] if pods else None desired = deployment.spec.replicas or 0 ready = False reason: str | None = None restarts = 0 updated_at = deployment.metadata.creation_timestamp if pod is not None: statuses = pod.status.container_statuses or [] ready = bool(statuses) and all(item.ready for item in statuses) restarts = sum(item.restart_count for item in statuses) for item in statuses: state = item.state if state and state.waiting and state.waiting.reason: reason = state.waiting.reason break if state and state.terminated and state.terminated.reason: reason = state.terminated.reason break updated_at = pod.status.start_time or pod.metadata.creation_timestamp if desired == 0: status = "stopped" elif reason in { "CrashLoopBackOff", "Error", "ImagePullBackOff", "ErrImagePull", "CreateContainerConfigError", "RunContainerError", } or (pod is not None and pod.status.phase == "Failed"): status = "error" elif ready and (deployment.status.available_replicas or 0) > 0: status = "running" else: status = "pending" reason = reason or (pod.status.phase if pod is not None else "Scheduling") container_spec = deployment.spec.template.spec.containers[0] limits = (container_spec.resources.limits or {}) if container_spec.resources else {} pvc_name = annotations.get(PVC_ANNOTATION) or self._deployment_pvc(deployment) storage = self._pvc_storage(namespace, pvc_name) cpu_usage, memory_usage = self._pod_metrics(namespace, pod.metadata.name if pod else None) return InstanceSummary( instance_id=instance_id, display_name=annotations.get(DISPLAY_ANNOTATION, instance_id), namespace=namespace, deployment=deployment.metadata.name, pod=pod.metadata.name if pod else None, status=status, reason=reason, ready=ready, restarts=restarts, pvc=pvc_name, storage=storage, image=container_spec.image, cpu_limit=limits.get("cpu"), memory_limit=limits.get("memory"), cpu_usage=cpu_usage, memory_usage=memory_usage, updated_at=_as_datetime(updated_at), legacy=legacy, deletable=not legacy and namespace == self.settings.namespace, ) def _pods_for_deployment(self, namespace: str, deployment: Any) -> list[Any]: selector = _selector(deployment.spec.selector.match_labels) try: pods = self.core.list_namespaced_pod( namespace, label_selector=selector, ).items except ApiException as exc: raise self._api_error(exc, "Could not list userbot Pods") from exc return sorted( pods, key=lambda pod: pod.metadata.creation_timestamp or datetime.min.replace(tzinfo=UTC), reverse=True, ) def _pod_metrics(self, namespace: str, pod_name: str | None) -> tuple[str | None, str | None]: if not pod_name: return None, None try: metrics = self.custom.get_namespaced_custom_object( "metrics.k8s.io", "v1beta1", namespace, "pods", pod_name, ) except (ApiException, AttributeError): return None, None containers = metrics.get("containers", []) cpu = containers[0].get("usage", {}).get("cpu") if containers else None memory = containers[0].get("usage", {}).get("memory") if containers else None return cpu, memory def _pvc_storage(self, namespace: str, pvc_name: str | None) -> str | None: if not pvc_name: return None try: pvc = self.core.read_namespaced_persistent_volume_claim(pvc_name, namespace) except ApiException: return None requests = pvc.spec.resources.requests or {} return requests.get("storage") @staticmethod def _deployment_pvc(deployment: Any) -> str | None: for volume in deployment.spec.template.spec.volumes or []: if volume.persistent_volume_claim: return volume.persistent_volume_claim.claim_name return None @staticmethod def _resource_names(instance_id: str) -> dict[str, str]: return { "deployment": f"userbot-{instance_id}", "secret": f"userbot-{instance_id}-credentials", "pvc": f"userbot-{instance_id}-data", } def _metadata( self, account: AccountBase, names: dict[str, str], resource_name: str, ) -> client.V1ObjectMeta: return client.V1ObjectMeta( name=resource_name, namespace=self.settings.namespace, labels={ "app.kubernetes.io/name": "userbot", INSTANCE_LABEL: account.instance_id, MANAGED_BY_LABEL: "userbot-panel", }, annotations={ DISPLAY_ANNOTATION: account.display_name, CREDENTIALS_ANNOTATION: names["secret"], PVC_ANNOTATION: names["pvc"], }, ) def _secret( self, account: AccountBase, session_string: str, names: dict[str, str], ) -> client.V1Secret: return client.V1Secret( metadata=self._metadata(account, names, names["secret"]), type="Opaque", string_data={ "API_ID": str(account.api_id), "API_HASH": account.api_hash, "STRINGSESSION": session_string, }, ) def _pvc( self, account: AccountBase, names: dict[str, str], ) -> client.V1PersistentVolumeClaim: return client.V1PersistentVolumeClaim( metadata=self._metadata(account, names, names["pvc"]), spec=client.V1PersistentVolumeClaimSpec( access_modes=["ReadWriteOnce"], storage_class_name=self.settings.storage_class, resources=client.V1VolumeResourceRequirements( requests={"storage": account.resources.storage}, ), ), ) def _deployment( self, account: AccountBase, names: dict[str, str], ) -> client.V1Deployment: pod_labels = { "app.kubernetes.io/name": "userbot", INSTANCE_LABEL: account.instance_id, MANAGED_BY_LABEL: "userbot-panel", } container = client.V1Container( name="userbot", image=self.settings.image, image_pull_policy="Always", env_from=[ client.V1EnvFromSource( secret_ref=client.V1SecretEnvSource(name=self.settings.common_secret) ), client.V1EnvFromSource( config_map_ref=client.V1ConfigMapEnvSource(name=self.settings.common_config) ), client.V1EnvFromSource( secret_ref=client.V1SecretEnvSource(name=names["secret"]) ), ], resources=client.V1ResourceRequirements( requests={"cpu": "80m", "memory": "512Mi"}, limits={ "cpu": account.resources.cpu_limit, "memory": account.resources.memory_limit, }, ), volume_mounts=[ client.V1VolumeMount(name="data", mount_path="/app/data"), client.V1VolumeMount(name="downloads", mount_path="/app/downloads"), ], ) pod_spec = client.V1PodSpec( service_account_name="userbot-runtime", automount_service_account_token=False, containers=[container], termination_grace_period_seconds=30, volumes=[ client.V1Volume( name="data", persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( claim_name=names["pvc"] ), ), client.V1Volume( name="downloads", host_path=client.V1HostPathVolumeSource( path=self.settings.downloads_host_path, type="DirectoryOrCreate", ), ), ], ) return client.V1Deployment( metadata=self._metadata(account, names, names["deployment"]), spec=client.V1DeploymentSpec( replicas=1, strategy=client.V1DeploymentStrategy(type="Recreate"), selector=client.V1LabelSelector(match_labels=pod_labels), template=client.V1PodTemplateSpec( metadata=client.V1ObjectMeta(labels=pod_labels), spec=pod_spec, ), ), ) def _rollback(self, created: list[tuple[str, str]]) -> None: for kind, name in reversed(created): try: if kind == "deployment": self.apps.delete_namespaced_deployment(name, self.settings.namespace) elif kind == "pvc": self.core.delete_namespaced_persistent_volume_claim( name, self.settings.namespace, ) else: self.core.delete_namespaced_secret(name, self.settings.namespace) except ApiException: pass @staticmethod def _api_error(exc: ApiException, detail: str) -> PanelError: if exc.status == 403: return PanelError(503, f"{detail}: Kubernetes RBAC denied the operation") if exc.status == 409: return PanelError(409, f"{detail}: resource conflict") if exc.status == 404: return PanelError(404, f"{detail}: resource not found") return PanelError(503, detail)