File: serve.py

package info (click to toggle)
pytest-httpbin 2.1.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 176 kB
  • sloc: python: 312; sh: 18; makefile: 14
file content (140 lines) | stat: -rw-r--r-- 4,236 bytes parent folder | download
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import os
import ssl
import threading
from urllib.parse import urljoin
from wsgiref.handlers import SimpleHandler
from wsgiref.simple_server import WSGIRequestHandler, WSGIServer, make_server

CERT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "certs")


class ServerHandler(SimpleHandler):
    server_software = "Pytest-HTTPBIN/0.1.0"
    http_version = "1.1"

    def cleanup_headers(self):
        SimpleHandler.cleanup_headers(self)
        self.headers["Connection"] = "Close"

    def close(self):
        try:
            self.request_handler.log_request(
                self.status.split(" ", 1)[0], self.bytes_sent
            )
        finally:
            SimpleHandler.close(self)


class Handler(WSGIRequestHandler):
    def handle(self):
        """Handle a single HTTP request"""

        self.raw_requestline = self.rfile.readline()
        if not self.parse_request():  # An error code has been sent, just exit
            return

        handler = ServerHandler(
            self.rfile, self.wfile, self.get_stderr(), self.get_environ()
        )
        handler.request_handler = self  # backpointer for logging
        handler.run(self.server.get_app())

    def get_environ(self):
        """
        wsgiref simple server adds content-type text/plain to everything, this
        removes it if it's not actually in the headers.
        """
        # Note: Can't use super since this is an oldstyle class in python 2.x
        environ = WSGIRequestHandler.get_environ(self).copy()
        if self.headers.get("content-type") is None:
            del environ["CONTENT_TYPE"]
        return environ


class SecureWSGIServer(WSGIServer):
    def get_request(self):
        socket, address = super().get_request()
        try:
            socket.settimeout(1.0)
            context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            context.load_cert_chain(
                os.path.join(CERT_DIR, "server.pem"),
                os.path.join(CERT_DIR, "server.key"),
            )
            return (
                context.wrap_socket(
                    socket, server_side=True, suppress_ragged_eofs=True
                ),
                address,
            )
        except Exception as e:
            print("pytest-httpbin server hit an exception serving request: %s" % e)
            print("attempting to ignore so the rest of the tests can run")
            raise

    def setup_environ(self):
        super().setup_environ()
        self.base_environ["HTTPS"] = "yes"


class Server:
    """
    HTTP server running a WSGI application in its own thread.
    """

    port_envvar = "HTTPBIN_HTTP_PORT"

    def __init__(self, host="127.0.0.1", port=0, application=None, **kwargs):
        self.app = application
        if self.port_envvar in os.environ:
            port = int(os.environ[self.port_envvar])
        self._server = make_server(
            host, port, self.app, handler_class=Handler, **kwargs
        )
        self.host = self._server.server_address[0]
        self.port = self._server.server_address[1]
        self.protocol = "http"

        self._thread = threading.Thread(
            name=self.__class__,
            target=self._server.serve_forever,
        )

    def __del__(self):
        if hasattr(self, "_server"):
            self.stop()

    def start(self):
        self._thread.start()

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, *args, **kwargs):
        self.stop()
        suppress_exc = self._server.__exit__(*args, **kwargs)
        self._thread.join()
        return suppress_exc

    def __add__(self, other):
        return self.url + other

    def stop(self):
        self._server.shutdown()

    @property
    def url(self):
        return f"{self.protocol}://{self.host}:{self.port}"

    def join(self, url, allow_fragments=True):
        return urljoin(self.url, url, allow_fragments=allow_fragments)


class SecureServer(Server):
    port_envvar = "HTTPBIN_HTTPS_PORT"

    def __init__(self, host="127.0.0.1", port=0, application=None, **kwargs):
        kwargs["server_class"] = SecureWSGIServer
        super().__init__(host, port, application, **kwargs)
        self.protocol = "https"