codekingpro/portable-devtools

1from __future__ import annotations
2
3import dataclasses
4import urllib.parse
5import urllib.request
6
7from .datastructures import Headers
8from .exceptions import InvalidProxy
9from .headers import build_authorization_basic, build_host
10from .http11 import USER_AGENT
11from .uri import DELIMS, WebSocketURI
12
13
14__all__ = ["get_proxy", "parse_proxy", "Proxy"]
15
16
17@dataclasses.dataclass
18class Proxy:
19    """
20    Proxy address.
21
22    Attributes:
23        scheme: ``"socks5h"``, ``"socks5"``, ``"socks4a"``, ``"socks4"``,
24            ``"https"``, or ``"http"``.
25        host: Normalized to lower case.
26        port: Always set even if it's the default.
27        username: Available when the proxy address contains `User Information`_.
28        password: Available when the proxy address contains `User Information`_.
29
30    .. _User Information: https://datatracker.ietf.org/doc/html/rfc3986#section-3.2.1
31
32    """
33
34    scheme: str
35    host: str
36    port: int
37    username: str | None = None
38    password: str | None = None
39
40    @property
41    def user_info(self) -> tuple[str, str] | None:
42        if self.username is None:
43            return None
44        assert self.password is not None
45        return (self.username, self.password)
46
47
48def parse_proxy(proxy: str) -> Proxy:
49    """
50    Parse and validate a proxy.
51
52    Args:
53        proxy: proxy.
54
55    Returns:
56        Parsed proxy.
57
58    Raises:
59        InvalidProxy: If ``proxy`` isn't a valid proxy.
60
61    """
62    parsed = urllib.parse.urlparse(proxy)
63    if parsed.scheme not in ["socks5h", "socks5", "socks4a", "socks4", "https", "http"]:
64        raise InvalidProxy(proxy, f"scheme {parsed.scheme} isn't supported")
65    if parsed.hostname is None:
66        raise InvalidProxy(proxy, "hostname isn't provided")
67    if parsed.path not in ["", "/"]:
68        raise InvalidProxy(proxy, "path is meaningless")
69    if parsed.query != "":
70        raise InvalidProxy(proxy, "query is meaningless")
71    if parsed.fragment != "":
72        raise InvalidProxy(proxy, "fragment is meaningless")
73
74    scheme = parsed.scheme
75    host = parsed.hostname
76    port = parsed.port or (443 if parsed.scheme == "https" else 80)
77    username = parsed.username
78    password = parsed.password
79    # urllib.parse.urlparse accepts URLs with a username but without a
80    # password. This doesn't make sense for HTTP Basic Auth credentials.
81    if username is not None and password is None:
82        raise InvalidProxy(proxy, "username provided without password")
83
84    try:
85        proxy.encode("ascii")
86    except UnicodeEncodeError:
87        # Input contains non-ASCII characters.
88        # It must be an IRI. Convert it to a URI.
89        host = host.encode("idna").decode()
90        if username is not None:
91            assert password is not None
92            username = urllib.parse.quote(username, safe=DELIMS)
93            password = urllib.parse.quote(password, safe=DELIMS)
94
95    return Proxy(scheme, host, port, username, password)
96
97
98def get_proxy(uri: WebSocketURI) -> str | None:
99    """
100    Return the proxy to use for connecting to the given WebSocket URI, if any.
101
102    """
103    if urllib.request.proxy_bypass(f"{uri.host}:{uri.port}"):
104        return None
105
106    # According to the _Proxy Usage_ section of RFC 6455, use a SOCKS5 proxy if
107    # available, else favor the proxy for HTTPS connections over the proxy for
108    # HTTP connections.
109
110    # The priority of a proxy for WebSocket connections is unspecified. We give
111    # it the highest priority. This makes it easy to configure a specific proxy
112    # for websockets.
113
114    # getproxies() may return SOCKS proxies as {"socks": "http://host:port"} or
115    # as {"https": "socks5h://host:port"} depending on whether they're declared
116    # in the operating system or in environment variables.
117
118    proxies = urllib.request.getproxies()
119    if uri.secure:
120        schemes = ["wss", "socks", "https"]
121    else:
122        schemes = ["ws", "socks", "https", "http"]
123
124    for scheme in schemes:
125        proxy = proxies.get(scheme)
126        if proxy is not None:
127            if scheme == "socks" and proxy.startswith("http://"):
128                proxy = "socks5h://" + proxy[7:]
129            return proxy
130    else:
131        return None
132
133
134def prepare_connect_request(
135    proxy: Proxy,
136    ws_uri: WebSocketURI,
137    user_agent_header: str | None = USER_AGENT,
138) -> bytes:
139    host = build_host(ws_uri.host, ws_uri.port, ws_uri.secure, always_include_port=True)
140    headers = Headers()
141    headers["Host"] = build_host(ws_uri.host, ws_uri.port, ws_uri.secure)
142    if user_agent_header is not None:
143        headers["User-Agent"] = user_agent_header
144    if proxy.username is not None:
145        assert proxy.password is not None  # enforced by parse_proxy()
146        headers["Proxy-Authorization"] = build_authorization_basic(
147            proxy.username, proxy.password
148        )
149    # We cannot use the Request class because it supports only GET requests.
150    return f"CONNECT {host} HTTP/1.1\r\n".encode() + headers.serialize()
151