summaryrefslogtreecommitdiffstats
path: root/ttun_server/endpoints.py
blob: eae0ebeeda059422b9a336ddbdc6094ab7fbcdc7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import logging
from base64 import b64decode, b64encode
from uuid import uuid4

from starlette.endpoints import HTTPEndpoint
from starlette.requests import Request
from starlette.responses import Response

from ttun_server.proxy_queue import ProxyQueue
from ttun_server.types import HttpRequestData, Message, HttpMessageType, HttpMessage

logger = logging.getLogger(__name__)


class HeaderMapping:
    def __init__(self, headers: list[tuple[str, str]]):
        self._headers = headers

    def items(self):
        for header in self._headers:
            yield header


class Proxy(HTTPEndpoint):
    async def dispatch(self) -> None:
        request = Request(self.scope, self.receive)

        [subdomain, *_] = request.headers['host'].split('.')
        response = Response(content='Not Found', status_code=404)

        identifier = str(uuid4())
        response_queue = await ProxyQueue.create_for_identifier(f'{subdomain}_{identifier}')

        try:

            request_queue = await ProxyQueue.get_for_identifier(subdomain)

            logger.debug('PROXY %s%s ', subdomain, request.url)
            await request_queue.enqueue(
                HttpMessage(
                    type=HttpMessageType.request.value,
                    identifier=identifier,
                    payload=
                    HttpRequestData(
                        method=request.method,
                        path=str(request.url).replace(str(request.base_url), '/'),
                        headers=list(request.headers.items()),
                        body=b64encode(await request.body()).decode()
                    )
                )
            )

            _response = await response_queue.dequeue()
            payload = _response['payload']
            response = Response(
                status_code=payload['status'],
                headers=HeaderMapping(payload['headers']),
                content=b64decode(payload['body'].encode())
            )
        except AssertionError:
            pass
        finally:
            await response(self.scope, self.receive, self.send)
            await response_queue.delete()


class Health(HTTPEndpoint):
    async def get(self, _) -> None:
        response = Response(content='OK', status_code=200)

        await response(self.scope, self.receive, self.send)