# Copyright (c) Microsoft Corporation. # # 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 abc import asyncio import contextlib import gzip import mimetypes import os import socket import threading from contextlib import closing from http import HTTPStatus from pathlib import Path from typing import ( Any, Callable, Dict, Generator, Generic, List, Optional, Set, Tuple, TypeVar, cast, ) from urllib.parse import urlparse from autobahn.twisted.resource import WebSocketResource from autobahn.twisted.websocket import WebSocketServerFactory, WebSocketServerProtocol from OpenSSL import crypto from twisted.internet import reactor as _twisted_reactor from twisted.internet import ssl from twisted.internet.selectreactor import SelectReactor from twisted.web import http _dirname = Path(os.path.abspath(__file__)).parent reactor = cast(SelectReactor, _twisted_reactor) def find_free_port() -> int: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) return s.getsockname()[1] T = TypeVar("T") class ExpectResponse(Generic[T]): def __init__(self) -> None: self._value: T @property def value(self) -> T: if not hasattr(self, "_value"): raise ValueError("no received value") return self._value def _write_error_body(request: "TestServerRequest", message: bytes) -> None: """Give error responses a Content-Type and a body. Gecko never fires `load` for an error response carrying neither -- the document renders through the plaintext viewer and then sits at readyState "interactive" forever. page.goto() waits for `load` by default, so it hangs until the test times out. Reproduced identically on stock Firefox, so this is Gecko behaviour rather than anything Camoufox does; it surfaced here only because this server was answering 404/401 with a bare status line, which no real server does. """ body = b"" + message + b"" request.setHeader(b"Content-Type", b"text/html; charset=utf-8") request.setHeader(b"Content-Length", str(len(body))) request.write(body) class TestServerRequest(http.Request): __test__ = False channel: "TestServerHTTPChannel" post_body: Optional[bytes] = None def process(self) -> None: server = self.channel.factory.server_instance if self.content: self.post_body = self.content.read() self.content.seek(0, 0) else: self.post_body = None uri = urlparse(self.uri.decode()) path = uri.path request_subscriber = server.request_subscribers.get(path) if request_subscriber: request_subscriber._loop.call_soon_threadsafe(request_subscriber.set_result, self) server.request_subscribers.pop(path) if path == "/ws": server._ws_resource.render(self) return if server.auth.get(path): authorization_header = self.requestHeaders.getRawHeaders("authorization") creds_correct = False if authorization_header: creds_correct = server.auth.get(path) == ( self.getUser().decode(), self.getPassword().decode(), ) if not creds_correct: self.setHeader(b"www-authenticate", 'Basic realm="Secure Area"') self.setResponseCode(HTTPStatus.UNAUTHORIZED) _write_error_body(self, b"401 Unauthorized") self.finish() return if server.csp.get(path): self.setHeader(b"Content-Security-Policy", server.csp[path]) if server.routes.get(path): server.routes[path](self) return file_content = None try: file_content = (server.static_path / path[1:]).read_bytes() content_type = mimetypes.guess_type(path)[0] if content_type and content_type.startswith("text/"): content_type += "; charset=utf-8" self.setHeader(b"Content-Type", content_type) self.setHeader(b"Cache-Control", "no-cache, no-store") if path in server.gzip_routes: self.setHeader("Content-Encoding", "gzip") self.write(gzip.compress(file_content)) else: self.setHeader(b"Content-Length", str(len(file_content))) self.write(file_content) self.setResponseCode(HTTPStatus.OK) except (FileNotFoundError, IsADirectoryError, PermissionError): self.setResponseCode(HTTPStatus.NOT_FOUND) _write_error_body(self, b"404 Not Found") self.finish() class TestServerHTTPChannel(http.HTTPChannel): factory: "TestServerFactory" requestFactory = TestServerRequest class TestServerFactory(http.HTTPFactory): server_instance: "Server" protocol = TestServerHTTPChannel class Server: protocol = "http" def __init__(self) -> None: self.PORT = find_free_port() self.EMPTY_PAGE = f"{self.protocol}://localhost:{self.PORT}/empty.html" self.PREFIX = f"{self.protocol}://localhost:{self.PORT}" self.CROSS_PROCESS_PREFIX = f"{self.protocol}://127.0.0.1:{self.PORT}" # On Windows, this list can be empty, reporting text/plain for scripts. mimetypes.add_type("text/html", ".html") mimetypes.add_type("text/css", ".css") mimetypes.add_type("application/javascript", ".js") mimetypes.add_type("image/png", ".png") mimetypes.add_type("font/woff2", ".woff2") def __repr__(self) -> str: return self.PREFIX @abc.abstractmethod def listen(self, factory: TestServerFactory) -> None: pass def start(self) -> None: request_subscribers: Dict[str, asyncio.Future] = {} auth: Dict[str, Tuple[str, str]] = {} csp: Dict[str, str] = {} routes: Dict[str, Callable[[TestServerRequest], Any]] = {} gzip_routes: Set[str] = set() self.request_subscribers = request_subscribers self.auth = auth self.csp = csp self.routes = routes self._ws_handlers: List[Callable[["WebSocketProtocol"], None]] = [] self.gzip_routes = gzip_routes self.static_path = _dirname / "assets" factory = TestServerFactory() factory.server_instance = self ws_factory = WebSocketServerFactory() ws_factory.protocol = WebSocketProtocol ws_factory.server_instance = self self._ws_resource = WebSocketResource(ws_factory) self.listen(factory) async def wait_for_request(self, path: str) -> TestServerRequest: if path in self.request_subscribers: return await self.request_subscribers[path] future: asyncio.Future["TestServerRequest"] = asyncio.Future() self.request_subscribers[path] = future return await future @contextlib.contextmanager def expect_request(self, path: str) -> Generator[ExpectResponse[TestServerRequest], None, None]: future = asyncio.create_task(self.wait_for_request(path)) cb_wrapper: ExpectResponse[TestServerRequest] = ExpectResponse() def done_cb(task: asyncio.Task) -> None: cb_wrapper._value = future.result() future.add_done_callback(done_cb) yield cb_wrapper def set_auth(self, path: str, username: str, password: str) -> None: self.auth[path] = (username, password) def set_csp(self, path: str, value: str) -> None: self.csp[path] = value def reset(self) -> None: self.request_subscribers.clear() self.auth.clear() self.csp.clear() self.gzip_routes.clear() self.routes.clear() self._ws_handlers.clear() def set_route(self, path: str, callback: Callable[[TestServerRequest], Any]) -> None: self.routes[path] = callback def enable_gzip(self, path: str) -> None: self.gzip_routes.add(path) def set_redirect(self, from_: str, to: str) -> None: def handle_redirect(request: http.Request) -> None: request.setResponseCode(HTTPStatus.FOUND) request.setHeader("location", to) request.finish() self.set_route(from_, handle_redirect) def send_on_web_socket_connection(self, data: bytes) -> None: self.once_web_socket_connection(lambda ws: ws.sendMessage(data)) def once_web_socket_connection(self, handler: Callable[["WebSocketProtocol"], None]) -> None: self._ws_handlers.append(handler) class HTTPServer(Server): def listen(self, factory: http.HTTPFactory) -> None: reactor.listenTCP(self.PORT, factory, interface="127.0.0.1") try: reactor.listenTCP(self.PORT, factory, interface="::1") except Exception: pass class HTTPSServer(Server): protocol = "https" def listen(self, factory: http.HTTPFactory) -> None: cert = ssl.PrivateCertificate.fromCertificateAndKeyPair( ssl.Certificate.loadPEM((_dirname / "testserver" / "cert.pem").read_bytes()), ssl.KeyPair.load( (_dirname / "testserver" / "key.pem").read_bytes(), crypto.FILETYPE_PEM ), ) contextFactory = cert.options() reactor.listenSSL(self.PORT, factory, contextFactory, interface="127.0.0.1") try: reactor.listenSSL(self.PORT, factory, contextFactory, interface="::1") except Exception: pass class WebSocketProtocol(WebSocketServerProtocol): def onOpen(self) -> None: for handler in self.factory.server_instance._ws_handlers.copy(): self.factory.server_instance._ws_handlers.remove(handler) handler(self) class TestServer: def __init__(self) -> None: self.server = HTTPServer() self.https_server = HTTPSServer() def start(self) -> None: self.server.start() self.https_server.start() self.thread = threading.Thread(target=lambda: reactor.run(installSignalHandlers=False)) self.thread.start() def stop(self) -> None: reactor.stop() self.thread.join() def reset(self) -> None: self.server.reset() self.https_server.reset() test_server = TestServer()