codekingpro/portable-devtools

1from time import monotonic2 3import requests4from urllib3.util import SKIP_HEADER, SKIPPABLE_HEADERS5 6from enum import Enum, auto7from typing import Iterable, Union, NamedTuple8from urllib.parse import urlsplit9 10from .cli.constants import (11    OUT_REQ_BODY,12    OUT_REQ_HEAD,13    OUT_RESP_BODY,14    OUT_RESP_HEAD,15    OUT_RESP_META16)17from .compat import cached_property18from .utils import split_cookies, parse_content_type_header19 20ELAPSED_TIME_LABEL = 'Elapsed time'21 22 23class HTTPMessage:24    """Abstract class for HTTP messages."""25 26    def __init__(self, orig):27        self._orig = orig28 29    def iter_body(self, chunk_size: int) -> Iterable[bytes]:30        """Return an iterator over the body."""31        raise NotImplementedError32 33    def iter_lines(self, chunk_size: int) -> Iterable[bytes]:34        """Return an iterator over the body yielding (`line`, `line_feed`)."""35        raise NotImplementedError36 37    @property38    def headers(self) -> str:39        """Return a `str` with the message's headers."""40        raise NotImplementedError41 42    @property43    def metadata(self) -> str:44        """Return metadata about the current message."""45        raise NotImplementedError46 47    @cached_property48    def encoding(self) -> str:49        ct, params = parse_content_type_header(self.content_type)50        return params.get('charset', '')51 52    @property53    def content_type(self) -> str:54        """Return the message content type."""55        ct = self._orig.headers.get('Content-Type', '')56        if not isinstance(ct, str):57            ct = ct.decode()58        return ct59 60 61class HTTPResponse(HTTPMessage):62    """A :class:`requests.models.Response` wrapper."""63 64    def iter_body(self, chunk_size=1):65        return self._orig.iter_content(chunk_size=chunk_size)66 67    def iter_lines(self, chunk_size):68        return ((line, b'\n') for line in self._orig.iter_lines(chunk_size))69 70    @property71    def headers(self):72        original = self._orig73        status_line = f'HTTP/{self.version} {original.status_code} {original.reason}'74        headers = [status_line]75        headers.extend(76            ': '.join(header)77            for header in original.headers.items()78            if header[0] != 'Set-Cookie'79        )80        headers.extend(81            f'Set-Cookie: {cookie}'82            for header, value in original.headers.items()83            for cookie in split_cookies(value)84            if header == 'Set-Cookie'85        )86        return '\r\n'.join(headers)87 88    @property89    def metadata(self) -> str:90        data = {}91        time_to_parse_headers = self._orig.elapsed.total_seconds()92        # noinspection PyProtectedMember93        time_since_headers_parsed = monotonic() - self._orig._httpie_headers_parsed_at94        time_elapsed = time_to_parse_headers + time_since_headers_parsed95        # data['Headers time'] = str(round(time_to_parse_headers, 5)) + 's'96        # data['Body time'] = str(round(time_since_headers_parsed, 5)) + 's'97        data[ELAPSED_TIME_LABEL] = str(round(time_elapsed, 10)) + 's'98        return '\n'.join(99            f'{key}: {value}'100            for key, value in data.items()101        )102 103    @property104    def version(self) -> str:105        """106        Return the HTTP version used by the server, e.g. '1.1'.107 108        Assume HTTP/1.1 if version is not available.109 110        """111        mapping = {112            9: '0.9',113            10: '1.0',114            11: '1.1',115            20: '2.0',116        }117        fallback = 11118        version = None119        try:120            raw = self._orig.raw121            if getattr(raw, '_original_response', None):122                version = raw._original_response.version123            else:124                version = raw.version125        except AttributeError:126            pass127        return mapping[version or fallback]128 129 130class HTTPRequest(HTTPMessage):131    """A :class:`requests.models.Request` wrapper."""132 133    def iter_body(self, chunk_size):134        yield self.body135 136    def iter_lines(self, chunk_size):137        yield self.body, b''138 139    @property140    def headers(self):141        url = urlsplit(self._orig.url)142 143        request_line = '{method} {path}{query} HTTP/1.1'.format(144            method=self._orig.method,145            path=url.path or '/',146            query=f'?{url.query}' if url.query else ''147        )148 149        headers = self._orig.headers.copy()150        if 'Host' not in self._orig.headers:151            headers['Host'] = url.netloc.split('@')[-1]152 153        headers = [154            f'{name}: {value if isinstance(value, str) else value.decode()}'155            for name, value in headers.items()156            if not (name.lower() in SKIPPABLE_HEADERS and value == SKIP_HEADER)157        ]158 159        headers.insert(0, request_line)160        headers = '\r\n'.join(headers).strip()161        return headers162 163    @property164    def body(self):165        body = self._orig.body166        if isinstance(body, str):167            # Happens with JSON/form request data parsed from the command line.168            body = body.encode()169        return body or b''170 171 172RequestsMessage = Union[requests.PreparedRequest, requests.Response]173 174 175class RequestsMessageKind(Enum):176    REQUEST = auto()177    RESPONSE = auto()178 179 180def infer_requests_message_kind(message: RequestsMessage) -> RequestsMessageKind:181    if isinstance(message, requests.PreparedRequest):182        return RequestsMessageKind.REQUEST183    elif isinstance(message, requests.Response):184        return RequestsMessageKind.RESPONSE185    else:186        raise TypeError(f"Unexpected message type: {type(message).__name__}")187 188 189OPTION_TO_PARAM = {190    RequestsMessageKind.REQUEST: {191        'headers': OUT_REQ_HEAD,192        'body': OUT_REQ_BODY,193    },194    RequestsMessageKind.RESPONSE: {195        'headers': OUT_RESP_HEAD,196        'body': OUT_RESP_BODY,197        'meta': OUT_RESP_META198    }199}200 201 202class OutputOptions(NamedTuple):203    kind: RequestsMessageKind204    headers: bool205    body: bool206    meta: bool = False207 208    def any(self):209        return (210            self.headers211            or self.body212            or self.meta213        )214 215    @classmethod216    def from_message(217        cls,218        message: RequestsMessage,219        raw_args: str = '',220        **kwargs221    ):222        kind = infer_requests_message_kind(message)223 224        options = {225            option: param in raw_args226            for option, param in OPTION_TO_PARAM[kind].items()227        }228        options.update(kwargs)229 230        return cls(231            kind=kind,232            **options233        )234