consumers.py

import json
import time
import six
import ssl
import aiohttp
import asyncio
from urllib.parse import urljoin
from django.conf import settings
from django.core.cache import cache

from asgiref.sync import sync_to_async

from urllib3.response import HTTPResponse
from kubernetes.client import Configuration, exceptions
from kubernetes.client.api import core_v1_api
from kubernetes.stream import stream
from kubernetes.stream.ws_client import STDOUT_CHANNEL, STDERR_CHANNEL, ERROR_CHANNEL

from channels.exceptions import DenyConnection
from channels.generic.http import AsyncHttpConsumer
from channels.generic.websocket import AsyncWebsocketConsumer

from .models.app import App
from .models.volume import Volume
from .permissions import has_app_permission


class AppPermChecker(object):
    timeout = 60 * 60

    def __init__(self, scope):
        self.scope = scope

    async def has_perm(self):
        if self.scope["user"] is None:
            return False, "user not login"
        app_id = self.scope["url_route"]["kwargs"]["id"]
        key = f"permission:user:{self.scope["user"].id}:app:{app_id}"
        permission = await cache.aget(key)
        if permission is None:
            try:
                app = await App.objects.aget(id=app_id)
                permission = await sync_to_async(has_app_permission)(
                    self.scope["user"], app, "GET")
                if permission[0]:
                    await cache.aset(key, permission, timeout=self.timeout)
            except App.DoesNotExist:
                permission = (False, "app not exists")
        return permission


class BaseAppConsumer(AsyncWebsocketConsumer):

    async def connect(self):
        app_perm_checker = AppPermChecker(self.scope)
        is_ok, message = await app_perm_checker.has_perm()
        if is_ok:
            await self.accept()
            self.id = self.scope["url_route"]["kwargs"]["id"]
        else:
            raise DenyConnection(message)


class BaseK8sConsumer(BaseAppConsumer):

    @property
    def kubernetes(self):
        with open('/var/run/secrets/kubernetes.io/serviceaccount/token') as token_file:
            token = token_file.read()
        config = Configuration(host=settings.SCHEDULER_URL)
        config.api_key = {"authorization": "Bearer " + token}
        config.verify_ssl = settings.K8S_API_VERIFY_TLS
        if config.verify_ssl:
            config.ssl_ca_cert = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
        Configuration.set_default(config)
        return core_v1_api.CoreV1Api()


class AppPodLogsConsumer(BaseK8sConsumer):

    async def connect(self):
        await super().connect()
        self.running = False
        self.response = None
        self.conneted = True
        self.buffer = b''
        self.delimiter = b"\r\n"
        self.pod_name = self.scope["url_route"]["kwargs"]["name"]

    def reader(self, sock):
        self.buffer += sock.read()
        try:
            while self.buffer and self.buffer.endswith(self.delimiter):
                index = self.buffer.index(self.delimiter)
                length = int(self.buffer[:index], base=16)
                if length == 0:
                    asyncio.create_task(self.close(code=1000))
                    break
                start_pos = index + len(self.delimiter)
                end_pos = start_pos + length + len(self.delimiter)
                asyncio.create_task(
                    self.send(bytes_data=self.buffer[start_pos:end_pos].strip(self.delimiter)))
                self.buffer = self.buffer[end_pos:]
        except BaseException:
            asyncio.create_task(self.close(code=1000))

    async def receive(self, text_data=None, bytes_data=None):
        if self.running:
            return
        self.running = True
        data = json.loads(text_data)
        args = (self.pod_name, self.id)
        lines = data.get("lines", 300)
        follow = data.get("follow", False)
        previous = data.get("previous", False)
        kwargs = {
            "follow": follow,
            "container": data.get("container", ""),
            "_preload_content": not follow,
            "previous": previous
        }
        if lines > 0:
            kwargs["tail_lines"] = lines
        self.response = await sync_to_async(self.kubernetes.read_namespaced_pod_log)(
            *args, **kwargs)
        if follow:
            loop = asyncio.get_event_loop()
            loop.add_reader(
                self.response.connection.sock, self.reader, self.response.connection.sock)
        else:
            asyncio.create_task(self.send(text_data=self.response))
            asyncio.create_task(self.close(code=1000))

    async def disconnect(self, close_code):
        if isinstance(self.response, HTTPResponse):
            loop = asyncio.get_event_loop()
            loop.remove_reader(self.response.connection.sock)
            await sync_to_async(self.response.close)()
        self.conneted = False


class AppPodExecConsumer(BaseK8sConsumer):

    async def connect(self):
        await super().connect()
        self.stream = None
        self.conneted = True
        self.pod_name = self.scope["url_route"]["kwargs"]["name"]

    async def send(self, data, channel=STDOUT_CHANNEL):
        channel_prefix = chr(channel)
        if data is None:
            return
        elif isinstance(data, bytes):
            channel_prefix = six.binary_type(channel_prefix, "ascii")
            await super().send(bytes_data=channel_prefix+data)
        elif isinstance(data, str):
            await super().send(text_data=channel_prefix+data)

    async def wait(self):
        future, loop = asyncio.Future(), asyncio.get_event_loop()
        loop.add_reader(self.stream.sock, future.set_result, None)
        future.add_done_callback(lambda f: loop.remove_reader(self.stream.sock))
        await future

    async def task(self):
        try:
            deadline = time.time() + settings.DRYCC_APP_POD_EXEC_TIMEOUT
            while self.stream.is_open() and self.conneted and time.time() < deadline:
                try:
                    await self.wait()
                    self.stream.update()
                    for channel in (ERROR_CHANNEL, STDOUT_CHANNEL, STDERR_CHANNEL):
                        if channel in self.stream._channels:
                            data = self.stream.read_channel(channel)
                            await self.send(data, channel)
                except ssl.SSLEOFError:
                    break
        except exceptions.ApiException as e:
            await self.send(str(e), STDERR_CHANNEL)
        finally:
            await self.close(code=1000)

    async def disconnect(self, close_code):
        if self.stream:
            await sync_to_async(self.stream.close)()
        self.conneted = False

    async def receive(self, text_data=None, bytes_data=None):
        if self.stream is None and text_data is not None:
            args = (self.kubernetes.connect_get_namespaced_pod_exec, self.pod_name, self.id)
            kwargs = json.loads(text_data)
            kwargs.update({"stderr": True, "stdout": True, "_preload_content": False})
            self.stream = await sync_to_async(stream)(*args, **kwargs)
            asyncio.create_task(self.task())
        elif self.stream is not None:
            data = text_data if text_data else bytes_data
            channel, data = ord(data[0]), data[1:]
            await sync_to_async(self.stream.write_channel)(channel, data)
        else:
            raise ValueError("This operation is not supported!")


class FilerProxyConsumer(AsyncHttpConsumer):
    from .middleware import ChannelOAuthMiddleware
    chunk_size = 64 * 1024
    middleware = ChannelOAuthMiddleware(None)
    SKIP_REQUEST_HEADERS = {
        'host', 'connection', 'keep-alive', 'proxy-connection', 'te', 'trailers', 'upgrade',
    }
    SKIP_RESPONSE_HEADERS = {
        'connection', 'keep-alive', 'proxy-authenticate', 'proxy-authorization', 'te', 'trailers',
        'transfer-encoding', 'upgrade',
    }

    async def handle(self, body: bytes):
        path = self.scope["url_route"]["kwargs"]["path"]
        client = await self._get_client(url_path=f"{self.scope["path"].removesuffix(path)}webdav/")
        if not client:
            await self.send_response(status=404, body=b'app or volume not found')
            return
        if path in ['_/ping', '_/bind']:  # need authentication
            await self._handle_controller_request(client, path)
            return
        elif not path.startswith('webdav/') or (filer := await client.info()) is None:
            await self.send_response(status=428, body=b'filer service unavailable')
            return
        filer_target_url = "{}?{}".format(
            urljoin(filer["endpoint"], self.scope["path"]), self.scope.get("query_string"))
        method = self.scope['method'].upper()
        headers = {
            name_bytes.decode('latin-1').lower(): value_bytes.decode('latin-1')
            for name_bytes, value_bytes in self.scope.get('headers', [])
            if name_bytes.decode('latin-1').lower() not in self.SKIP_REQUEST_HEADERS
        }
        await self._handle_proxy_request(filer_target_url, method, headers, body)

    async def _get_client(self, url_path):
        try:
            from .filer import FilerClient
            app = await App.objects.aget(id=self.scope["url_route"]["kwargs"]["id"])
            volume = await Volume.objects.filter(
                app=app, name=self.scope["url_route"]["kwargs"]["name"]).afirst()
            if not volume:
                return None
            return FilerClient(app.id, volume, url_path)
        except App.DoesNotExist:
            return None

    async def _forward_response(self, response: aiohttp.ClientResponse):
        response_headers = [
            [name.encode('latin-1'), value.encode('latin-1')]
            for name, value in response.headers.items()
            if name.lower() not in self.SKIP_RESPONSE_HEADERS
        ]
        await self.send_headers(status=response.status, headers=response_headers)
        async for chunk in response.content.iter_chunked(self.chunk_size):
            if chunk:
                await self.send_body(chunk, more_body=True)
        # Send final empty chunk to indicate end of body
        await self.send_body(b'', more_body=False)

    async def _handle_proxy_request(self, url, method: str, headers: dict[str, str], data: bytes):
        async with aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=settings.DRYCC_FILER_DURATION)
        ) as session:
            try:
                async with session.request(
                    method=method, url=url, headers=headers, data=data, allow_redirects=False
                ) as response:
                    await self._forward_response(response)
            except aiohttp.ClientError as e:
                await self.send_response(502, f"proxy service unavailable: {e}".encode('utf-8'))
            except asyncio.TimeoutError:
                await self.send_response(504, b'proxy request to backend filer timeout')

    async def _handle_controller_request(self, client, path: str):
        await self.middleware.login(self.scope)
        app_perm_checker = AppPermChecker(self.scope)
        is_ok, message = await app_perm_checker.has_perm()
        status, body = 200, b''
        if not is_ok:
            status, body = 403, message.encode('utf-8')
        elif path == '_/ping':
            if (filer := await client.info()) is not None:
                body = b'pong'
            else:
                status, body = 503, b'filer service unavailable'
        elif path == '_/bind':
            filer = await client.bind()
            body = json.dumps({
                "username": filer["username"], "password": filer["password"],
            }).encode('utf-8')
        await self.send_response(status, body)