1import argparse
2import os
3import platform
4import sys
5from typing import List, Optional, Tuple, Union
6
7import requests
8from pygments import __version__ as pygments_version
9from requests import __version__ as requests_version
10
11from . import __version__ as httpie_version
12from .cli.constants import OUT_REQ_BODY, OUT_REQ_HEAD, OUT_RESP_BODY, OUT_RESP_HEAD
13from .client import collect_messages
14from .context import Environment
15from .downloads import Downloader
16from .output.writer import write_message, write_stream, MESSAGE_SEPARATOR_BYTES
17from .plugins.registry import plugin_manager
18from .status import ExitStatus, http_status_to_exit_status
19
20
21# noinspection PyDefaultArgument
22def main(args: List[Union[str, bytes]] = sys.argv, env=Environment()) -> ExitStatus:
23    """
24    The main function.
25
26    Pre-process args, handle some special types of invocations,
27    and run the main program with error handling.
28
29    Return exit status code.
30
31    """
32    program_name, *args = args
33    env.program_name = os.path.basename(program_name)
34    args = decode_raw_args(args, env.stdin_encoding)
35    plugin_manager.load_installed_plugins()
36
37    from .cli.definition import parser
38
39    if env.config.default_options:
40        args = env.config.default_options + args
41
42    include_debug_info = '--debug' in args
43    include_traceback = include_debug_info or '--traceback' in args
44
45    if include_debug_info:
46        print_debug_info(env)
47        if args == ['--debug']:
48            return ExitStatus.SUCCESS
49
50    exit_status = ExitStatus.SUCCESS
51
52    try:
53        parsed_args = parser.parse_args(
54            args=args,
55            env=env,
56        )
57    except KeyboardInterrupt:
58        env.stderr.write('\n')
59        if include_traceback:
60            raise
61        exit_status = ExitStatus.ERROR_CTRL_C
62    except SystemExit as e:
63        if e.code != ExitStatus.SUCCESS:
64            env.stderr.write('\n')
65            if include_traceback:
66                raise
67            exit_status = ExitStatus.ERROR
68    else:
69        try:
70            exit_status = program(
71                args=parsed_args,
72                env=env,
73            )
74        except KeyboardInterrupt:
75            env.stderr.write('\n')
76            if include_traceback:
77                raise
78            exit_status = ExitStatus.ERROR_CTRL_C
79        except SystemExit as e:
80            if e.code != ExitStatus.SUCCESS:
81                env.stderr.write('\n')
82                if include_traceback:
83                    raise
84                exit_status = ExitStatus.ERROR
85        except requests.Timeout:
86            exit_status = ExitStatus.ERROR_TIMEOUT
87            env.log_error(f'Request timed out ({parsed_args.timeout}s).')
88        except requests.TooManyRedirects:
89            exit_status = ExitStatus.ERROR_TOO_MANY_REDIRECTS
90            env.log_error(
91                f'Too many redirects'
92                f' (--max-redirects={parsed_args.max_redirects}).'
93            )
94        except Exception as e:
95            # TODO: Further distinction between expected and unexpected errors.
96            msg = str(e)
97            if hasattr(e, 'request'):
98                request = e.request
99                if hasattr(request, 'url'):
100                    msg = (
101                        f'{msg} while doing a {request.method}'
102                        f' request to URL: {request.url}'
103                    )
104            env.log_error(f'{type(e).__name__}: {msg}')
105            if include_traceback:
106                raise
107            exit_status = ExitStatus.ERROR
108
109    return exit_status
110
111
112def get_output_options(
113    args: argparse.Namespace,
114    message: Union[requests.PreparedRequest, requests.Response]
115) -> Tuple[bool, bool]:
116    return {
117        requests.PreparedRequest: (
118            OUT_REQ_HEAD in args.output_options,
119            OUT_REQ_BODY in args.output_options,
120        ),
121        requests.Response: (
122            OUT_RESP_HEAD in args.output_options,
123            OUT_RESP_BODY in args.output_options,
124        ),
125    }[type(message)]
126
127
128def program(args: argparse.Namespace, env: Environment) -> ExitStatus:
129    """
130    The main program without error handling.
131
132    """
133    # TODO: Refactor and drastically simplify, especially so that the separator logic is elsewhere.
134    exit_status = ExitStatus.SUCCESS
135    downloader = None
136    initial_request: Optional[requests.PreparedRequest] = None
137    final_response: Optional[requests.Response] = None
138
139    def separate():
140        getattr(env.stdout, 'buffer', env.stdout).write(MESSAGE_SEPARATOR_BYTES)
141
142    def request_body_read_callback(chunk: bytes):
143        should_pipe_to_stdout = bool(
144            # Request body output desired
145            OUT_REQ_BODY in args.output_options
146            # & not `.read()` already pre-request (e.g., for  compression)
147            and initial_request
148            # & non-EOF chunk
149            and chunk
150        )
151        if should_pipe_to_stdout:
152            msg = requests.PreparedRequest()
153            msg.is_body_upload_chunk = True
154            msg.body = chunk
155            msg.headers = initial_request.headers
156            write_message(requests_message=msg, env=env, args=args, with_body=True, with_headers=False)
157
158    try:
159        if args.download:
160            args.follow = True  # --download implies --follow.
161            downloader = Downloader(output_file=args.output_file, progress_file=env.stderr, resume=args.download_resume)
162            downloader.pre_request(args.headers)
163        messages = collect_messages(args=args, config_dir=env.config.directory,
164                                    request_body_read_callback=request_body_read_callback)
165        force_separator = False
166        prev_with_body = False
167
168        # Process messages as they’re generated
169        for message in messages:
170            is_request = isinstance(message, requests.PreparedRequest)
171            with_headers, with_body = get_output_options(args=args, message=message)
172            do_write_body = with_body
173            if prev_with_body and (with_headers or with_body) and (force_separator or not env.stdout_isatty):
174                # Separate after a previous message with body, if needed. See test_tokens.py.
175                separate()
176            force_separator = False
177            if is_request:
178                if not initial_request:
179                    initial_request = message
180                if with_body:
181                    is_streamed_upload = not isinstance(message.body, (str, bytes))
182                    do_write_body = not is_streamed_upload
183                    force_separator = is_streamed_upload and env.stdout_isatty
184            else:
185                final_response = message
186                if args.check_status or downloader:
187                    exit_status = http_status_to_exit_status(http_status=message.status_code, follow=args.follow)
188                    if exit_status != ExitStatus.SUCCESS and (not env.stdout_isatty or args.quiet == 1):
189                        env.log_error(f'HTTP {message.raw.status} {message.raw.reason}', level='warning')
190            write_message(requests_message=message, env=env, args=args, with_headers=with_headers,
191                          with_body=do_write_body)
192            prev_with_body = with_body
193
194        # Cleanup
195        if force_separator:
196            separate()
197        if downloader and exit_status == ExitStatus.SUCCESS:
198            # Last response body download.
199            download_stream, download_to = downloader.start(
200                initial_url=initial_request.url,
201                final_response=final_response,
202            )
203            write_stream(stream=download_stream, outfile=download_to, flush=False)
204            downloader.finish()
205            if downloader.interrupted:
206                exit_status = ExitStatus.ERROR
207                env.log_error(
208                    f'Incomplete download: size={downloader.status.total_size};'
209                    f' downloaded={downloader.status.downloaded}'
210                )
211        return exit_status
212
213    finally:
214        if downloader and not downloader.finished:
215            downloader.failed()
216        if args.output_file and args.output_file_specified:
217            args.output_file.close()
218
219
220def print_debug_info(env: Environment):
221    env.stderr.writelines([
222        f'HTTPie {httpie_version}\n',
223        f'Requests {requests_version}\n',
224        f'Pygments {pygments_version}\n',
225        f'Python {sys.version}\n{sys.executable}\n',
226        f'{platform.system()} {platform.release()}',
227    ])
228    env.stderr.write('\n\n')
229    env.stderr.write(repr(env))
230    env.stderr.write('\n\n')
231    env.stderr.write(repr(plugin_manager))
232    env.stderr.write('\n')
233
234
235def decode_raw_args(
236    args: List[Union[str, bytes]],
237    stdin_encoding: str
238) -> List[str]:
239    """
240    Convert all bytes args to str
241    by decoding them using stdin encoding.
242
243    """
244    return [
245        arg.decode(stdin_encoding)
246        if type(arg) is bytes else arg
247        for arg in args
248    ]
249