diff --git a/examples_asyncio/pod_exec.py b/examples_asyncio/pod_exec.py new file mode 100644 index 0000000000..36c52a92e0 --- /dev/null +++ b/examples_asyncio/pod_exec.py @@ -0,0 +1,149 @@ +import asyncio + +from aiohttp.http import WSMsgType + +from kubernetes.aio import client, config, utils +from kubernetes.aio.client.api_client import ApiClient +from kubernetes.aio.stream import WsApiClient +from kubernetes.aio.stream.ws_client import ( + ERROR_CHANNEL, STDERR_CHANNEL, STDOUT_CHANNEL, +) + +BUSYBOX_POD = "busybox-test" + + +async def find_busybox_pod(): + async with ApiClient() as api: + v1 = client.CoreV1Api(api) + ret = await v1.list_pod_for_all_namespaces() + for i in ret.items: + if i.metadata.namespace == 'default' and i.metadata.name == BUSYBOX_POD: + print(f"Found busybox pod: {i.metadata.name}") + return i.metadata.name + return None + + +async def create_busybox_pod(): + print(f"Pod {BUSYBOX_POD} does not exist. Creating it...") + manifest = { + 'apiVersion': 'v1', + 'kind': 'Pod', + 'metadata': { + 'name': BUSYBOX_POD, + }, + 'spec': { + 'containers': [{ + 'image': 'busybox', + 'name': 'sleep', + "args": [ + "/bin/sh", + "-c", + "while true; do date; sleep 5; done" + ] + }] + } + } + async with ApiClient() as api: + objects = await utils.create_from_dict(api, manifest, namespace="default") + pod = objects[0] + print(f"Created pod {pod.metadata.name}.") + return pod.metadata.name + + +async def wait_busybox_pod_ready(): + print(f"Waiting pod {BUSYBOX_POD} to be ready.") + async with ApiClient() as api: + v1 = client.CoreV1Api(api) + while True: + ret = await v1.read_namespaced_pod(name=BUSYBOX_POD, namespace="default") + if ret.status.phase != 'Pending': + break + await asyncio.sleep(1) + + +async def main(): + # Configs can be set in Configuration class directly or using helper + # utility. If no argument provided, the config will be loaded from + # default location. + await config.load_kube_config() + + pod = await find_busybox_pod() + if not pod: + pod = await create_busybox_pod() + await wait_busybox_pod_ready() + + # Execute a command in a pod non-interactively, and display its output + print("-------------") + async with WsApiClient() as ws_api: + v1_ws = client.CoreV1Api(api_client=ws_api) + exec_command = [ + "/bin/sh", + "-c", + "echo This message goes to stderr >&2; echo This message goes to stdout", + ] + ret = await v1_ws.connect_get_namespaced_pod_exec( + pod, + "default", + command=exec_command, + stderr=True, + stdin=False, + stdout=True, + tty=False, + ) + print(f"Response: {ret}") + + # Execute a command interactively. If _preload_content=False is passed to + # connect_get_namespaced_pod_exec(), the returned object is an aiohttp ClientWebSocketResponse + # object, that can be manipulated directly. + print("-------------") + async with WsApiClient() as ws_api: + v1_ws = client.CoreV1Api(api_client=ws_api) + exec_command = ['/bin/sh'] + websocket = await v1_ws.connect_get_namespaced_pod_exec( + BUSYBOX_POD, + "default", + command=exec_command, + stderr=True, + stdin=True, + stdout=True, + tty=False, + _preload_content=False, + ) + commands = [ + "echo 'This message goes to stdout'\n", + "echo 'This message goes to stderr' >&2\n", + "exit 1\n", + ] + error_data = "" + closed = False + async with websocket as ws: + while commands and not closed: + command = commands.pop(0) + stdin_channel_prefix = chr(0) + await ws.send_bytes((stdin_channel_prefix + command).encode("utf-8")) + while True: + try: + msg = await ws.receive(timeout=1) + except asyncio.TimeoutError: + break + if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): + closed = True + break + channel = msg.data[0] + data = msg.data[1:].decode("utf-8") + if not data: + continue + if channel == STDOUT_CHANNEL: + print(f"stdout: {data}") + elif channel == STDERR_CHANNEL: + print(f"stderr: {data}") + elif channel == ERROR_CHANNEL: + error_data += data + if error_data: + returncode = ws_api.parse_error_data(error_data) + print(f"Exit code: {returncode}") + +if __name__ == "__main__": + loop = asyncio.get_event_loop() + loop.run_until_complete(main()) + loop.close() diff --git a/kubernetes/aio/__init__.py b/kubernetes/aio/__init__.py index 2c4c1d0b85..97c5f2e383 100644 --- a/kubernetes/aio/__init__.py +++ b/kubernetes/aio/__init__.py @@ -19,7 +19,8 @@ import kubernetes.aio.client as client import kubernetes.aio.config as config import kubernetes.aio.dynamic as dynamic +import kubernetes.aio.stream as stream import kubernetes.aio.utils as utils import kubernetes.aio.watch as watch -__all__ = ["client", "config", "dynamic", "utils", "watch"] +__all__ = ["client", "config", "dynamic", "stream", "utils", "watch"] diff --git a/kubernetes/aio/stream/__init__.py b/kubernetes/aio/stream/__init__.py new file mode 100644 index 0000000000..8a75e3742d --- /dev/null +++ b/kubernetes/aio/stream/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2017 The Kubernetes Authors. +# +# 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. + +from .ws_client import WsApiClient diff --git a/kubernetes/aio/stream/ws_client.py b/kubernetes/aio/stream/ws_client.py new file mode 100644 index 0000000000..bcc0669e27 --- /dev/null +++ b/kubernetes/aio/stream/ws_client.py @@ -0,0 +1,111 @@ +# 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 json + +from six.moves.urllib.parse import urlencode, urlparse, urlunparse + +from kubernetes.aio.client import ApiClient +from kubernetes.aio.client.rest import RESTResponse + +STDIN_CHANNEL = 0 +STDOUT_CHANNEL = 1 +STDERR_CHANNEL = 2 +ERROR_CHANNEL = 3 +RESIZE_CHANNEL = 4 + + +def get_websocket_url(url): + parsed_url = urlparse(url) + parts = list(parsed_url) + if parsed_url.scheme == 'http': + parts[0] = 'ws' + elif parsed_url.scheme == 'https': + parts[0] = 'wss' + return urlunparse(parts) + + +class WsResponse(RESTResponse): + + def __init__(self, status, data): + self.status = status + self.data = data + self.headers = {} + self.reason = None + + def getheaders(self): + return self.headers + + def getheader(self, name, default=None): + return self.headers.get(name, default) + + +class WsApiClient(ApiClient): + + def __init__(self, configuration=None, header_name=None, header_value=None, + cookie=None, pool_threads=1, heartbeat=None): + super().__init__(configuration, header_name, header_value, cookie, pool_threads) + self.heartbeat = heartbeat + + @classmethod + def parse_error_data(cls, error_data): + """ + Parse data received on ERROR_CHANNEL and return the command exit code. + """ + error_data_json = json.loads(error_data) + if error_data_json.get("status") == "Success": + return 0 + return int(error_data_json["details"]["causes"][0]['message']) + + async def request(self, method, url, query_params=None, headers=None, + post_params=None, body=None, _preload_content=True, + _request_timeout=None): + + # Expand command parameter list to indivitual command params + if query_params: + new_query_params = [] + for key, value in query_params: + if key == 'command' and isinstance(value, list): + for command in value: + new_query_params.append((key, command)) + else: + new_query_params.append((key, value)) + query_params = new_query_params + + if headers is None: + headers = {} + if 'sec-websocket-protocol' not in headers: + headers['sec-websocket-protocol'] = 'v4.channel.k8s.io' + + if query_params: + url += '?' + urlencode(query_params) + + url = get_websocket_url(url) + + if _preload_content: + + resp_all = '' + async with self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat) as ws: + async for msg in ws: + msg = msg.data.decode('utf-8') + if len(msg) > 1: + channel = ord(msg[0]) + data = msg[1:] + if data: + if channel in [STDOUT_CHANNEL, STDERR_CHANNEL]: + resp_all += data + + return WsResponse(200, resp_all.encode('utf-8')) + + else: + + return self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat) diff --git a/kubernetes/aio/stream/ws_client_test.py b/kubernetes/aio/stream/ws_client_test.py new file mode 100644 index 0000000000..92dd41f1a6 --- /dev/null +++ b/kubernetes/aio/stream/ws_client_test.py @@ -0,0 +1,124 @@ +# Copyright 2017 The Kubernetes Authors. +# +# 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. + +from unittest import IsolatedAsyncioTestCase +from unittest.mock import Mock, patch + +from kubernetes.aio import client +from kubernetes.aio.stream import WsApiClient +from kubernetes.aio.stream.ws_client import WsResponse, get_websocket_url + + +class WsMock: + def __init__(self): + self.iter = 0 + + def __aiter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return self + + async def __aenter__(self): + return self + + async def __anext__(self): + self.iter += 1 + if self.iter > 5: + raise StopAsyncIteration + return WsResponse(200, (chr(1) + 'mock').encode('utf-8')) + + +class WSClientTest(IsolatedAsyncioTestCase): + + def test_websocket_client(self): + for url, ws_url in [ + ('http://localhost/api', 'ws://localhost/api'), + ('https://localhost/api', 'wss://localhost/api'), + ('https://domain.com/api', 'wss://domain.com/api'), + ('https://api.domain.com/api', 'wss://api.domain.com/api'), + ('http://api.domain.com', 'ws://api.domain.com'), + ('https://api.domain.com', 'wss://api.domain.com'), + ('http://api.domain.com/', 'ws://api.domain.com/'), + ('https://api.domain.com/', 'wss://api.domain.com/'), + ]: + self.assertEqual(get_websocket_url(url), ws_url) + + async def test_exec_ws(self): + mock = Mock() + mock.RESTClientObject.return_value.pool_manager = mock + mock.ws_connect.return_value = WsMock() + with patch('kubernetes.aio.client.api_client.rest', mock): + + api_client = WsApiClient() + api_client.configuration.host = 'https://localhost' + ws = client.CoreV1Api(api_client=api_client) + resp = ws.connect_get_namespaced_pod_exec('pod', 'namespace', + command="mock-command", + stderr=True, stdin=False, + stdout=True, tty=False) + + ret = await resp + self.assertEqual(ret, 'mock' * 5) + mock.ws_connect.assert_called_once_with( + 'wss://localhost/api/v1/namespaces/namespace/pods/pod/exec?' + 'command=mock-command&stderr=True&stdin=False&stdout=True&tty=False', + headers={ + 'sec-websocket-protocol': 'v4.channel.k8s.io', + 'Accept': '*/*', + 'User-Agent': api_client.user_agent + }, + heartbeat=None + ) + + async def test_exec_ws_with_heartbeat(self): + mock = Mock() + mock.RESTClientObject.return_value.pool_manager = mock + mock.ws_connect.return_value = WsMock() + with patch('kubernetes.aio.client.api_client.rest', mock): + + api_client = WsApiClient(heartbeat=30) + api_client.configuration.host = 'https://localhost' + ws = client.CoreV1Api(api_client=api_client) + resp = ws.connect_get_namespaced_pod_exec('pod', 'namespace', + command='mock-command', + stderr=True, stdin=False, + stdout=True, tty=False) + + ret = await resp + self.assertEqual(ret, 'mock' * 5) + mock.ws_connect.assert_called_once_with( + 'wss://localhost/api/v1/namespaces/namespace/pods/pod/exec?' + 'command=mock-command&stderr=True&stdin=False&stdout=True&tty=False', + headers={ + 'sec-websocket-protocol': 'v4.channel.k8s.io', + 'Accept': '*/*', + 'User-Agent': api_client.user_agent + }, + heartbeat=30 + ) + + def test_parse_error_data_success(self): + error_data = '{"metadata":{},"status":"Success"}' + return_code = WsApiClient.parse_error_data(error_data) + self.assertEqual(return_code, 0) + + def test_parse_error_data_failure(self): + error_data = ( + '{"metadata":{},"status":"Failure",' + '"message":"command terminated with non-zero exit code",' + '"reason":"NonZeroExitCode",' + '"details":{"causes":[{"reason":"ExitCode","message":"1"}]}}') + return_code = WsApiClient.parse_error_data(error_data) + self.assertEqual(return_code, 1)