codekingpro/portable-devtools

1from __future__ import annotations
2
3import http
4import ssl as ssl_module
5import urllib.parse
6from typing import Any, Callable, Literal
7
8from ..http11 import Request, Response
9from .server import Server, ServerConnection, serve
10
11
12__all__ = ["route", "unix_route", "Router"]
13
14
15try:
16    from werkzeug.exceptions import NotFound
17    from werkzeug.routing import Map, RequestRedirect
18
19except ImportError:
20
21    def route(
22        url_map: Map,
23        *args: Any,
24        server_name: str | None = None,
25        ssl: ssl_module.SSLContext | Literal[True] | None = None,
26        create_router: type[Router] | None = None,
27        **kwargs: Any,
28    ) -> Server:
29        raise ImportError("route() requires werkzeug")
30
31    def unix_route(
32        url_map: Map,
33        path: str | None = None,
34        **kwargs: Any,
35    ) -> Server:
36        raise ImportError("unix_route() requires werkzeug")
37
38else:
39
40    def route(
41        url_map: Map,
42        *args: Any,
43        server_name: str | None = None,
44        ssl: ssl_module.SSLContext | Literal[True] | None = None,
45        create_router: type[Router] | None = None,
46        **kwargs: Any,
47    ) -> Server:
48        """
49        Create a WebSocket server dispatching connections to different handlers.
50
51        This feature requires the third-party library `werkzeug`_:
52
53        .. code-block:: console
54
55            $ pip install werkzeug
56
57        .. _werkzeug: https://werkzeug.palletsprojects.com/
58
59        :func:`route` accepts the same arguments as
60        :func:`~websockets.sync.server.serve`, except as described below.
61
62        The first argument is a :class:`werkzeug.routing.Map` that maps URL patterns
63        to connection handlers. In addition to the connection, handlers receive
64        parameters captured in the URL as keyword arguments.
65
66        Here's an example::
67
68
69            from websockets.sync.router import route
70            from werkzeug.routing import Map, Rule
71
72            def channel_handler(websocket, channel_id):
73                ...
74
75            url_map = Map([
76                Rule("/channel/<uuid:channel_id>", endpoint=channel_handler),
77                ...
78            ])
79
80            with route(url_map, ...) as server:
81                server.serve_forever()
82
83        Refer to the documentation of :mod:`werkzeug.routing` for details.
84
85        If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map,
86        when the server runs behind a reverse proxy that modifies the ``Host``
87        header or terminates TLS, you need additional configuration:
88
89        * Set ``server_name`` to the name of the server as seen by clients. When
90          not provided, websockets uses the value of the ``Host`` header.
91
92        * Set ``ssl=True`` to generate ``wss://`` URIs without enabling TLS.
93          Under the hood, this bind the URL map with a ``url_scheme`` of
94          ``wss://`` instead of ``ws://``.
95
96        There is no need to specify ``websocket=True`` in each rule. It is added
97        automatically.
98
99        Args:
100            url_map: Mapping of URL patterns to connection handlers.
101            server_name: Name of the server as seen by clients. If :obj:`None`,
102                websockets uses the value of the ``Host`` header.
103            ssl: Configuration for enabling TLS on the connection. Set it to
104                :obj:`True` if a reverse proxy terminates TLS connections.
105            create_router: Factory for the :class:`Router` dispatching requests to
106                handlers. Set it to a wrapper or a subclass to customize routing.
107
108        """
109        url_scheme = "ws" if ssl is None else "wss"
110        if ssl is not True and ssl is not None:
111            kwargs["ssl"] = ssl
112
113        if create_router is None:
114            create_router = Router
115
116        router = create_router(url_map, server_name, url_scheme)
117
118        _process_request: (
119            Callable[
120                [ServerConnection, Request],
121                Response | None,
122            ]
123            | None
124        ) = kwargs.pop("process_request", None)
125        if _process_request is None:
126            process_request: Callable[
127                [ServerConnection, Request],
128                Response | None,
129            ] = router.route_request
130        else:
131
132            def process_request(
133                connection: ServerConnection, request: Request
134            ) -> Response | None:
135                response = _process_request(connection, request)
136                if response is not None:
137                    return response
138                return router.route_request(connection, request)
139
140        return serve(router.handler, *args, process_request=process_request, **kwargs)
141
142    def unix_route(
143        url_map: Map,
144        path: str | None = None,
145        **kwargs: Any,
146    ) -> Server:
147        """
148        Create a WebSocket Unix server dispatching connections to different handlers.
149
150        :func:`unix_route` combines the behaviors of :func:`route` and
151        :func:`~websockets.sync.server.unix_serve`.
152
153        Args:
154            url_map: Mapping of URL patterns to connection handlers.
155            path: File system path to the Unix socket.
156
157        """
158        return route(url_map, unix=True, path=path, **kwargs)
159
160
161class Router:
162    """WebSocket router supporting :func:`route`."""
163
164    def __init__(
165        self,
166        url_map: Map,
167        server_name: str | None = None,
168        url_scheme: str = "ws",
169    ) -> None:
170        self.url_map = url_map
171        self.server_name = server_name
172        self.url_scheme = url_scheme
173        for rule in self.url_map.iter_rules():
174            rule.websocket = True
175
176    def get_server_name(self, connection: ServerConnection, request: Request) -> str:
177        if self.server_name is None:
178            return request.headers["Host"]
179        else:
180            return self.server_name
181
182    def redirect(self, connection: ServerConnection, url: str) -> Response:
183        response = connection.respond(http.HTTPStatus.FOUND, f"Found at {url}")
184        response.headers["Location"] = url
185        return response
186
187    def not_found(self, connection: ServerConnection) -> Response:
188        return connection.respond(http.HTTPStatus.NOT_FOUND, "Not Found")
189
190    def route_request(
191        self, connection: ServerConnection, request: Request
192    ) -> Response | None:
193        """Route incoming request."""
194        url_map_adapter = self.url_map.bind(
195            server_name=self.get_server_name(connection, request),
196            url_scheme=self.url_scheme,
197        )
198        try:
199            parsed = urllib.parse.urlparse(request.path)
200            handler, kwargs = url_map_adapter.match(
201                path_info=parsed.path,
202                query_args=parsed.query,
203            )
204        except RequestRedirect as redirect:
205            return self.redirect(connection, redirect.new_url)
206        except NotFound:
207            return self.not_found(connection)
208        connection.handler, connection.handler_kwargs = handler, kwargs
209        return None
210
211    def handler(self, connection: ServerConnection) -> None:
212        """Handle a connection."""
213        return connection.handler(connection, **connection.handler_kwargs)
214