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