Coverage for scripts / example_job_validation / kube.py: 100.00%
134 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Cluster access for live harnesses: SSM tunnel plus kubectl execution.
3The historical example harness keeps its default behavior when no explicit
4kubeconfig is supplied. Sibling harnesses can instead pass an isolated path;
5AWS CLI and kubectl receive ``--kubeconfig`` and every nested ``gco`` process
6receives ``KUBECONFIG``, so those runs never rewrite ``~/.kube/config``.
7"""
9from __future__ import annotations
11import os
12import stat
13import subprocess
14import time
15from collections.abc import Callable, Iterator, Mapping
16from contextlib import contextmanager
17from pathlib import Path
18from typing import Any
20import yaml
22#: How the harness invokes kubectl; a function taking kubectl args and
23#: returning (exit_code, stdout, stderr).
24KubectlRunner = Callable[..., tuple[int, str, str]]
26_CLUSTER_API_READY_TIMEOUT_SECONDS = 45.0
27_CLUSTER_API_PROBE_TIMEOUT_SECONDS = 8
28_CLUSTER_API_RETRY_SECONDS = 1.0
29_PERMANENT_API_STARTUP_MARKERS = (
30 "certificate signed by unknown authority",
31 "error loading config file",
32 "exec plugin: invalid apiversion",
33 "forbidden",
34 "invalid configuration",
35 "no configuration has been provided",
36 "the server has asked for the client to provide credentials",
37 "tls: failed to verify certificate",
38 "unauthorized",
39 "x509:",
40)
43def _is_permanent_api_startup_error(detail: str) -> bool:
44 normalized = detail.casefold()
45 return any(marker in normalized for marker in _PERMANENT_API_STARTUP_MARKERS)
48def _wait_for_cluster_api(
49 kubectl: KubectlRunner,
50 *,
51 tunnel_process: subprocess.Popen[bytes] | None,
52 timeout_seconds: float = _CLUSTER_API_READY_TIMEOUT_SECONDS,
53 poll_interval_seconds: float = _CLUSTER_API_RETRY_SECONDS,
54) -> None:
55 """Wait until the tunnel can carry an authenticated Kubernetes API request."""
56 if timeout_seconds <= 0:
57 raise ValueError("timeout_seconds must be positive")
58 if poll_interval_seconds <= 0:
59 raise ValueError("poll_interval_seconds must be positive")
61 from cli import ssm_tunnel
63 deadline = time.monotonic() + timeout_seconds
64 attempts = 0
65 last_error = "probe was not attempted"
66 while True:
67 if tunnel_process is not None and (
68 detail := ssm_tunnel.exited_api_tunnel_detail(tunnel_process)
69 ):
70 raise RuntimeError(
71 "SSM tunnel exited before the Kubernetes API became ready: " + detail
72 )
74 remaining = deadline - time.monotonic()
75 if attempts and remaining <= 0:
76 raise RuntimeError(
77 "Kubernetes API did not become ready through the SSM tunnel within "
78 f"{timeout_seconds:.1f}s after {attempts} attempt(s). "
79 f"Last transient error: {last_error[:1000]}"
80 )
81 command_timeout = max(
82 1,
83 min(_CLUSTER_API_PROBE_TIMEOUT_SECONDS, int(max(remaining, 0.0)) + 1),
84 )
85 request_timeout = min(command_timeout, 5)
86 attempts += 1
87 try:
88 returncode, stdout, stderr = kubectl(
89 f"--request-timeout={request_timeout}s",
90 "get",
91 "--raw=/readyz",
92 timeout=command_timeout,
93 )
94 except subprocess.TimeoutExpired as exc:
95 detail = f"kubectl readiness probe timed out after {exc.timeout}s"
96 else:
97 if returncode == 0:
98 return
99 detail = (stderr or stdout).strip() or f"kubectl exited with status {returncode}"
100 if _is_permanent_api_startup_error(detail):
101 raise RuntimeError(
102 "Kubernetes API readiness probe failed with a permanent error: " + detail[:1000]
103 )
104 last_error = detail
106 remaining = deadline - time.monotonic()
107 if remaining <= 0:
108 raise RuntimeError(
109 "Kubernetes API did not become ready through the SSM tunnel within "
110 f"{timeout_seconds:.1f}s after {attempts} attempt(s). "
111 f"Last transient error: {last_error[:1000]}"
112 )
113 time.sleep(min(poll_interval_seconds, remaining))
116class _QuietFormatter:
117 """Adapter for cli formatter callbacks used by the tunnel helpers."""
119 @staticmethod
120 def print_info(message: str) -> None:
121 print(f"[tunnel] {message}")
123 print_success = print_info
124 print_warning = print_info
125 print_error = print_info
128def _kubeconfig_path(kubeconfig_path: Path | None = None) -> Path:
129 return kubeconfig_path if kubeconfig_path is not None else Path.home() / ".kube" / "config"
132def _environment_with_kubeconfig(
133 kubeconfig_path: Path | None,
134 base: Mapping[str, str] | None = None,
135) -> dict[str, str] | None:
136 if kubeconfig_path is None:
137 return dict(base) if base is not None else None
138 environment = dict(base) if base is not None else dict(os.environ)
139 environment["KUBECONFIG"] = str(kubeconfig_path)
140 return environment
143def _update_kubeconfig_command(
144 cluster_name: str,
145 region: str,
146 kubeconfig_path: Path | None,
147) -> list[str]:
148 command = ["aws", "eks", "update-kubeconfig", "--name", cluster_name, "--region", region]
149 if kubeconfig_path is not None:
150 command.extend(("--kubeconfig", str(kubeconfig_path)))
151 return command
154def _validate_and_secure_isolated_kubeconfig(path: Path) -> None:
155 """Require the AWS-written isolated kubeconfig to be a current-user regular file."""
156 metadata = path.lstat()
157 if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode):
158 raise ValueError(f"Isolated kubeconfig must be a regular file: {path}")
159 if hasattr(os, "geteuid") and metadata.st_uid != os.geteuid():
160 raise PermissionError(f"Isolated kubeconfig is not owned by this user: {path}")
161 if os.name != "nt":
162 path.chmod(0o600)
165def refresh_kubeconfig(
166 cluster_name: str,
167 region: str,
168 *,
169 kubeconfig_path: Path | None = None,
170) -> Path:
171 """Run AWS CLI update-kubeconfig against the selected config file."""
172 subprocess.run(
173 _update_kubeconfig_command(cluster_name, region, kubeconfig_path),
174 check=True,
175 capture_output=True,
176 text=True,
177 env=_environment_with_kubeconfig(kubeconfig_path),
178 shell=False,
179 )
180 path = _kubeconfig_path(kubeconfig_path)
181 if kubeconfig_path is not None:
182 _validate_and_secure_isolated_kubeconfig(path)
183 return path
186def update_and_point_kubeconfig_at_tunnel(
187 cluster_name: str,
188 region: str,
189 server: str,
190 tls_server_name: str,
191 *,
192 kubeconfig_path: Path | None = None,
193) -> None:
194 """Refresh kubeconfig, point it at the tunnel, and preserve real TLS SNI."""
195 path = refresh_kubeconfig(
196 cluster_name,
197 region,
198 kubeconfig_path=kubeconfig_path,
199 )
200 config = yaml.safe_load(path.read_text(encoding="utf-8"))
201 if not isinstance(config, dict) or not isinstance(config.get("clusters"), list):
202 raise ValueError(f"Kubeconfig has no cluster list: {path}")
203 expected_suffix = f"cluster/{cluster_name}"
204 matched = False
205 for entry in config["clusters"]:
206 if not isinstance(entry, dict) or not str(entry.get("name", "")).endswith(expected_suffix):
207 continue
208 cluster = entry.get("cluster")
209 if not isinstance(cluster, dict):
210 raise ValueError(f"Kubeconfig cluster entry is malformed: {path}")
211 cluster["server"] = server
212 cluster["tls-server-name"] = tls_server_name
213 matched = True
214 if not matched:
215 raise ValueError(f"Kubeconfig did not contain the requested cluster: {cluster_name}")
216 path.write_text(yaml.safe_dump(config), encoding="utf-8")
217 if kubeconfig_path is not None:
218 _validate_and_secure_isolated_kubeconfig(path)
221def ensure_cluster_access_entry(
222 repo_root: Path,
223 region: str,
224 *,
225 kubeconfig_path: Path | None = None,
226 gco_command: tuple[str, ...] = ("gco",),
227) -> None:
228 """Grant cluster-admin through an explicitly selected GCO checkout."""
229 if not gco_command or any(not isinstance(part, str) or not part for part in gco_command):
230 raise ValueError("gco_command must be a non-empty argv prefix")
231 result = subprocess.run(
232 [*gco_command, "stacks", "access", "--region", region],
233 cwd=repo_root,
234 capture_output=True,
235 text=True,
236 env=_environment_with_kubeconfig(kubeconfig_path),
237 shell=False,
238 )
239 if result.returncode != 0:
240 raise RuntimeError(f"gco stacks access failed for {region}: {result.stderr.strip()[:500]}")
243@contextmanager
244def cluster_session(
245 repo_root: Path,
246 cluster_name: str,
247 region: str,
248 *,
249 kubeconfig_path: Path | None = None,
250 gco_command: tuple[str, ...] = ("gco",),
251) -> Iterator[KubectlRunner]:
252 """Access entry plus tunnel for one region; optionally isolate kubeconfig."""
253 from cli import cluster_tunnel
255 formatter = _QuietFormatter()
256 ensure_cluster_access_entry(
257 repo_root,
258 region,
259 kubeconfig_path=kubeconfig_path,
260 gco_command=gco_command,
261 )
262 with cluster_tunnel.open_api_server_tunnel(
263 formatter,
264 cluster=cluster_name,
265 region=region,
266 via_ssm=cluster_tunnel.AUTO_BASTION,
267 assume_yes=True,
268 ) as session:
269 if session.active and session.server and session.tls_server_name:
270 update_and_point_kubeconfig_at_tunnel(
271 cluster_name,
272 region,
273 session.server,
274 session.tls_server_name,
275 kubeconfig_path=kubeconfig_path,
276 )
277 else:
278 refresh_kubeconfig(
279 cluster_name,
280 region,
281 kubeconfig_path=kubeconfig_path,
282 )
284 def kubectl(*args: str, timeout: int = 120, **kwargs: Any) -> tuple[int, str, str]:
285 if kwargs.pop("shell", False):
286 raise ValueError("cluster_session kubectl does not allow shell execution")
287 command = ["kubectl"]
288 if kubeconfig_path is not None:
289 command.extend(("--kubeconfig", str(kubeconfig_path)))
290 command.extend(args)
291 caller_environment = kwargs.pop("env", None)
292 environment = _environment_with_kubeconfig(
293 kubeconfig_path,
294 caller_environment,
295 )
296 result = subprocess.run(
297 command,
298 capture_output=True,
299 text=True,
300 timeout=timeout,
301 env=environment,
302 shell=False,
303 **kwargs,
304 )
305 return result.returncode, result.stdout, result.stderr
307 if session.active:
308 _wait_for_cluster_api(
309 kubectl,
310 tunnel_process=getattr(session, "process", None),
311 )
312 yield kubectl