|
| 1 | +#!/usr/bin/env python3 |
| 2 | +import logging |
| 3 | +from pathlib import Path |
| 4 | +from typing import Optional |
| 5 | + |
| 6 | +import click |
| 7 | +import yaml |
| 8 | +from kubernetes.stream import stream |
| 9 | + |
| 10 | +from warnet.constants import HookValue |
| 11 | +from warnet.k8s import ( |
| 12 | + get_default_namespace, |
| 13 | + get_static_client, |
| 14 | +) |
| 15 | +from warnet.process import run_command |
| 16 | + |
| 17 | +MISSION = "hello" |
| 18 | +PRIMARY_CONTAINER = MISSION |
| 19 | + |
| 20 | +PLUGIN_DIR_TAG = "plugin_dir" |
| 21 | + |
| 22 | + |
| 23 | +class PluginError(Exception): |
| 24 | + pass |
| 25 | + |
| 26 | + |
| 27 | +log = logging.getLogger(MISSION) |
| 28 | +if not log.hasHandlers(): |
| 29 | + console_handler = logging.StreamHandler() |
| 30 | + console_handler.setLevel(logging.DEBUG) |
| 31 | + formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") |
| 32 | + console_handler.setFormatter(formatter) |
| 33 | + log.addHandler(console_handler) |
| 34 | +log.setLevel(logging.DEBUG) |
| 35 | +log.propagate = True |
| 36 | + |
| 37 | + |
| 38 | +@click.group() |
| 39 | +@click.pass_context |
| 40 | +def hello(ctx): |
| 41 | + """Commands for the Hello plugin""" |
| 42 | + ctx.ensure_object(dict) |
| 43 | + plugin_dir = Path(__file__).resolve().parent |
| 44 | + ctx.obj[PLUGIN_DIR_TAG] = Path(plugin_dir) |
| 45 | + |
| 46 | + |
| 47 | +@hello.command() |
| 48 | +@click.argument("network_file_path", type=str) |
| 49 | +@click.argument("hook_value", type=str) |
| 50 | +@click.argument("namespace", type=str) |
| 51 | +@click.argument("nargs", nargs=-1) |
| 52 | +@click.pass_context |
| 53 | +def entrypoint(ctx, network_file_path: str, hook_value: str, namespace: str, nargs): |
| 54 | + """Plugin entrypoint""" |
| 55 | + assert hook_value in { |
| 56 | + item.value for item in HookValue |
| 57 | + }, f"{hook_value} is not a valid HookValue" |
| 58 | + |
| 59 | + network_file_path = Path(network_file_path) |
| 60 | + |
| 61 | + with network_file_path.open() as f: |
| 62 | + network_file = yaml.safe_load(f) or {} |
| 63 | + if not isinstance(network_file, dict): |
| 64 | + raise ValueError(f"Invalid network file structure: {network_file_path}") |
| 65 | + |
| 66 | + plugins_section = network_file.get("plugins", {}) |
| 67 | + hook_section = plugins_section.get(hook_value, {}) |
| 68 | + |
| 69 | + plugin_name = Path(__file__).resolve().parent.stem |
| 70 | + plugin_data = hook_section.get(plugin_name) |
| 71 | + if not plugin_data: |
| 72 | + raise PluginError(f"Could not find {plugin_name} in {network_file_path}") |
| 73 | + |
| 74 | + _entrypoint(ctx, plugin_data, HookValue(hook_value), namespace, nargs) |
| 75 | + |
| 76 | + |
| 77 | +def _entrypoint(ctx, plugin_data: dict, hook_value: HookValue, namespace: str, nargs): |
| 78 | + """Called by entrypoint""" |
| 79 | + match hook_value: |
| 80 | + case ( |
| 81 | + HookValue.PRE_NETWORK |
| 82 | + | HookValue.POST_NETWORK |
| 83 | + | HookValue.PRE_DEPLOY |
| 84 | + | HookValue.POST_DEPLOY |
| 85 | + ): |
| 86 | + data = get_data(plugin_data) |
| 87 | + if data: |
| 88 | + _launch_pod(ctx, install_name=hook_value.value.lower() + "-hello", **data) |
| 89 | + else: |
| 90 | + _launch_pod(ctx, install_name=hook_value.value.lower() + "-hello") |
| 91 | + case HookValue.PRE_NODE: |
| 92 | + name = nargs[0] + "-pre-hello-pod" |
| 93 | + _launch_pod(ctx, install_name=hook_value.value.lower() + "-" + name, podName=name) |
| 94 | + case HookValue.POST_NODE: |
| 95 | + name = nargs[0] + "-post-hello-pod" |
| 96 | + _launch_pod(ctx, install_name=hook_value.value.lower() + "-" + name, podName=name) |
| 97 | + |
| 98 | + |
| 99 | +def get_data(plugin_data: dict) -> Optional[dict]: |
| 100 | + data = {key: plugin_data.get(key) for key in ("podName", "helloTo") if plugin_data.get(key)} |
| 101 | + return data or None |
| 102 | + |
| 103 | + |
| 104 | +def _launch_pod( |
| 105 | + ctx, install_name: str = "hello", podName: str = "hello-pod", helloTo: str = "World!" |
| 106 | +): |
| 107 | + command = f"helm upgrade --install {install_name} {ctx.obj[PLUGIN_DIR_TAG]}/charts/hello --set podName={podName} --set helloTo={helloTo}" |
| 108 | + log.info(command) |
| 109 | + log.info(run_command(command)) |
| 110 | + |
| 111 | + |
| 112 | +def _sh(pod, method: str, params: tuple[str, ...]) -> str: |
| 113 | + namespace = get_default_namespace() |
| 114 | + |
| 115 | + sclient = get_static_client() |
| 116 | + if params: |
| 117 | + cmd = [method] |
| 118 | + cmd.extend(params) |
| 119 | + else: |
| 120 | + cmd = [method] |
| 121 | + try: |
| 122 | + resp = stream( |
| 123 | + sclient.connect_get_namespaced_pod_exec, |
| 124 | + pod, |
| 125 | + namespace, |
| 126 | + container=PRIMARY_CONTAINER, |
| 127 | + command=cmd, |
| 128 | + stderr=True, |
| 129 | + stdin=False, |
| 130 | + stdout=True, |
| 131 | + tty=False, |
| 132 | + _preload_content=False, |
| 133 | + ) |
| 134 | + stdout = "" |
| 135 | + stderr = "" |
| 136 | + while resp.is_open(): |
| 137 | + resp.update(timeout=1) |
| 138 | + if resp.peek_stdout(): |
| 139 | + stdout_chunk = resp.read_stdout() |
| 140 | + stdout += stdout_chunk |
| 141 | + if resp.peek_stderr(): |
| 142 | + stderr_chunk = resp.read_stderr() |
| 143 | + stderr += stderr_chunk |
| 144 | + return stdout + stderr |
| 145 | + except Exception as err: |
| 146 | + print(f"Could not execute stream: {err}") |
| 147 | + |
| 148 | + |
| 149 | +@hello.command(context_settings={"ignore_unknown_options": True}) |
| 150 | +@click.argument("pod", type=str) |
| 151 | +@click.argument("method", type=str) |
| 152 | +@click.argument("params", type=str, nargs=-1) # this will capture all remaining arguments |
| 153 | +def sh(pod: str, method: str, params: tuple[str, ...]): |
| 154 | + """Run shell commands in a pod""" |
| 155 | + print(_sh(pod, method, params)) |
| 156 | + |
| 157 | + |
| 158 | +if __name__ == "__main__": |
| 159 | + hello() |
0 commit comments