1from __future__ import absolute_import
2
3from base64 import b64encode
4
5from ..exceptions import UnrewindableBodyError
6from ..packages.six import b, integer_types
7
8# Pass as a value within ``headers`` to skip
9# emitting some HTTP headers that are added automatically.
10# The only headers that are supported are ``Accept-Encoding``,
11# ``Host``, and ``User-Agent``.
12SKIP_HEADER = "@@@SKIP_HEADER@@@"
13SKIPPABLE_HEADERS = frozenset(["accept-encoding", "host", "user-agent"])
14
15ACCEPT_ENCODING = "gzip,deflate"
16try:
17    import brotli as _unused_module_brotli  # noqa: F401
18except ImportError:
19    pass
20else:
21    ACCEPT_ENCODING += ",br"
22
23_FAILEDTELL = object()
24
25
26def make_headers(
27    keep_alive=None,
28    accept_encoding=None,
29    user_agent=None,
30    basic_auth=None,
31    proxy_basic_auth=None,
32    disable_cache=None,
33):
34    """
35    Shortcuts for generating request headers.
36
37    :param keep_alive:
38        If ``True``, adds 'connection: keep-alive' header.
39
40    :param accept_encoding:
41        Can be a boolean, list, or string.
42        ``True`` translates to 'gzip,deflate'.
43        List will get joined by comma.
44        String will be used as provided.
45
46    :param user_agent:
47        String representing the user-agent you want, such as
48        "python-urllib3/0.6"
49
50    :param basic_auth:
51        Colon-separated username:password string for 'authorization: basic ...'
52        auth header.
53
54    :param proxy_basic_auth:
55        Colon-separated username:password string for 'proxy-authorization: basic ...'
56        auth header.
57
58    :param disable_cache:
59        If ``True``, adds 'cache-control: no-cache' header.
60
61    Example::
62
63        >>> make_headers(keep_alive=True, user_agent="Batman/1.0")
64        {'connection': 'keep-alive', 'user-agent': 'Batman/1.0'}
65        >>> make_headers(accept_encoding=True)
66        {'accept-encoding': 'gzip,deflate'}
67    """
68    headers = {}
69    if accept_encoding:
70        if isinstance(accept_encoding, str):
71            pass
72        elif isinstance(accept_encoding, list):
73            accept_encoding = ",".join(accept_encoding)
74        else:
75            accept_encoding = ACCEPT_ENCODING
76        headers["accept-encoding"] = accept_encoding
77
78    if user_agent:
79        headers["user-agent"] = user_agent
80
81    if keep_alive:
82        headers["connection"] = "keep-alive"
83
84    if basic_auth:
85        headers["authorization"] = "Basic " + b64encode(b(basic_auth)).decode("utf-8")
86
87    if proxy_basic_auth:
88        headers["proxy-authorization"] = "Basic " + b64encode(
89            b(proxy_basic_auth)
90        ).decode("utf-8")
91
92    if disable_cache:
93        headers["cache-control"] = "no-cache"
94
95    return headers
96
97
98def set_file_position(body, pos):
99    """
100    If a position is provided, move file to that point.
101    Otherwise, we'll attempt to record a position for future use.
102    """
103    if pos is not None:
104        rewind_body(body, pos)
105    elif getattr(body, "tell", None) is not None:
106        try:
107            pos = body.tell()
108        except (IOError, OSError):
109            # This differentiates from None, allowing us to catch
110            # a failed `tell()` later when trying to rewind the body.
111            pos = _FAILEDTELL
112
113    return pos
114
115
116def rewind_body(body, body_pos):
117    """
118    Attempt to rewind body to a certain position.
119    Primarily used for request redirects and retries.
120
121    :param body:
122        File-like object that supports seek.
123
124    :param int pos:
125        Position to seek to in file.
126    """
127    body_seek = getattr(body, "seek", None)
128    if body_seek is not None and isinstance(body_pos, integer_types):
129        try:
130            body_seek(body_pos)
131        except (IOError, OSError):
132            raise UnrewindableBodyError(
133                "An error occurred when rewinding request body for redirect/retry."
134            )
135    elif body_pos is _FAILEDTELL:
136        raise UnrewindableBodyError(
137            "Unable to record file position for rewinding "
138            "request body during a redirect/retry."
139        )
140    else:
141        raise ValueError(
142            "body_pos must be of type integer, instead it was %s." % type(body_pos)
143        )
144