codekingpro/portable-devtools

1"""2Persistent, JSON-serialized sessions.3 4"""5import os6import re7 8from http.cookies import SimpleCookie9from http.cookiejar import Cookie10from pathlib import Path11from typing import Any, Dict, List, Optional, Union12 13from requests.auth import AuthBase14from requests.cookies import RequestsCookieJar, remove_cookie_by_name15 16from .context import Environment, LogLevel17from .cookies import HTTPieCookiePolicy18from .cli.dicts import HTTPHeadersDict19from .config import BaseConfigDict, DEFAULT_CONFIG_DIR20from .utils import url_as_host21from .plugins.registry import plugin_manager22 23from .legacy import (24    v3_1_0_session_cookie_format as legacy_cookies,25    v3_2_0_session_header_format as legacy_headers26)27 28 29SESSIONS_DIR_NAME = 'sessions'30DEFAULT_SESSIONS_DIR = DEFAULT_CONFIG_DIR / SESSIONS_DIR_NAME31VALID_SESSION_NAME_PATTERN = re.compile('^[a-zA-Z0-9_.-]+$')32# Request headers starting with these prefixes won't be stored in sessions.33# They are specific to each request.34# <https://en.wikipedia.org/wiki/List_of_HTTP_header_fields#Requests>35SESSION_IGNORED_HEADER_PREFIXES = ['Content-', 'If-']36 37# Cookie related options38KEPT_COOKIE_OPTIONS = ['name', 'expires', 'path', 'value', 'domain', 'secure']39DEFAULT_COOKIE_PATH = '/'40 41 42def is_anonymous_session(session_name: str) -> bool:43    return os.path.sep in session_name44 45 46def session_hostname_to_dirname(hostname: str, session_name: str) -> str:47    # host:port => host_port48    hostname = hostname.replace(':', '_')49    return os.path.join(50        SESSIONS_DIR_NAME,51        hostname,52        f'{session_name}.json'53    )54 55 56def strip_port(hostname: str) -> str:57    return hostname.split(':')[0]58 59 60def materialize_cookie(cookie: Cookie) -> Dict[str, Any]:61    materialized_cookie = {62        option: getattr(cookie, option)63        for option in KEPT_COOKIE_OPTIONS64    }65 66    if (67        cookie._rest.get('is_explicit_none')68        and materialized_cookie['domain'] == ''69    ):70        materialized_cookie['domain'] = None71 72    return materialized_cookie73 74 75def materialize_cookies(jar: RequestsCookieJar) -> List[Dict[str, Any]]:76    return [77        materialize_cookie(cookie)78        for cookie in jar79    ]80 81 82def materialize_headers(headers: Dict[str, str]) -> List[Dict[str, Any]]:83    return [84        {85            'name': name,86            'value': value87        }88        for name, value in headers.copy().items()89    ]90 91 92def get_httpie_session(93    env: Environment,94    config_dir: Path,95    session_name: str,96    host: Optional[str],97    url: str,98    *,99    suppress_legacy_warnings: bool = False100) -> 'Session':101    bound_hostname = host or url_as_host(url)102    if not bound_hostname:103        # HACK/FIXME: httpie-unixsocket's URLs have no hostname.104        bound_hostname = 'localhost'105 106    if is_anonymous_session(session_name):107        path = os.path.expanduser(session_name)108        session_id = path109    else:110        path = config_dir / session_hostname_to_dirname(bound_hostname, session_name)111        session_id = session_name112 113    session = Session(114        path,115        env=env,116        session_id=session_id,117        bound_host=strip_port(bound_hostname),118        suppress_legacy_warnings=suppress_legacy_warnings119    )120    session.load()121    return session122 123 124class Session(BaseConfigDict):125    helpurl = 'https://httpie.io/docs#sessions'126    about = 'HTTPie session file'127 128    def __init__(129        self,130        path: Union[str, Path],131        env: Environment,132        bound_host: str,133        session_id: str,134        suppress_legacy_warnings: bool = False,135    ):136        super().__init__(path=Path(path))137 138        # Default values for the session files139        self['headers'] = []140        self['cookies'] = []141        self['auth'] = {142            'type': None,143            'username': None,144            'password': None145        }146 147        # Runtime state of the Session objects.148        self.env = env149        self._headers = HTTPHeadersDict()150        self.cookie_jar = RequestsCookieJar(151            # See also a temporary workaround for a Requests bug in `compat.py`.152            policy=HTTPieCookiePolicy(),153        )154        self.session_id = session_id155        self.bound_host = bound_host156        self.suppress_legacy_warnings = suppress_legacy_warnings157 158    def _add_cookies(self, cookies: List[Dict[str, Any]]) -> None:159        for cookie in cookies:160            domain = cookie.get('domain', '')161            if domain is None:162                # domain = None means explicitly lack of cookie, though163                # requests requires domain to be a string so we'll cast it164                # manually.165                cookie['domain'] = ''166                cookie['rest'] = {'is_explicit_none': True}167 168            self.cookie_jar.set(**cookie)169 170    def pre_process_data(self, data: Dict[str, Any]) -> Dict[str, Any]:171        for key, deserializer, importer in [172            ('cookies', legacy_cookies.pre_process, self._add_cookies),173            ('headers', legacy_headers.pre_process, self._headers.update),174        ]:175            values = data.get(key)176            if values:177                normalized_values = deserializer(self, values)178            else:179                normalized_values = []180 181            importer(normalized_values)182 183        return data184 185    def post_process_data(self, data: Dict[str, Any]) -> Dict[str, Any]:186        for key, store, serializer, exporter in [187            ('cookies', self.cookie_jar, materialize_cookies, legacy_cookies.post_process),188            ('headers', self._headers, materialize_headers, legacy_headers.post_process),189        ]:190            original_type = type(data.get(key))191            values = serializer(store)192 193            data[key] = exporter(194                values,195                original_type=original_type196            )197 198        return data199 200    def _compute_new_headers(self, request_headers: HTTPHeadersDict) -> HTTPHeadersDict:201        new_headers = HTTPHeadersDict()202        for name, value in request_headers.copy().items():203            if value is None:204                continue  # Ignore explicitly unset headers205 206            original_value = value207            if type(value) is not str:208                value = value.decode()209 210            if name.lower() == 'user-agent' and value.startswith('HTTPie/'):211                continue212 213            if name.lower() == 'cookie':214                for cookie_name, morsel in SimpleCookie(value).items():215                    if not morsel['path']:216                        morsel['path'] = DEFAULT_COOKIE_PATH217                    self.cookie_jar.set(cookie_name, morsel)218 219                request_headers.remove_item(name, original_value)220                continue221 222            for prefix in SESSION_IGNORED_HEADER_PREFIXES:223                if name.lower().startswith(prefix.lower()):224                    break225            else:226                new_headers.add(name, value)227 228        return new_headers229 230    def update_headers(self, request_headers: HTTPHeadersDict):231        """232        Update the session headers with the request ones while ignoring233        certain name prefixes.234 235        """236 237        new_headers = self._compute_new_headers(request_headers)238        new_keys = new_headers.copy().keys()239 240        # New headers will take priority over the existing ones, and override241        # them directly instead of extending them.242        for key, value in self._headers.copy().items():243            if key in new_keys:244                continue245 246            new_headers.add(key, value)247 248        self._headers = new_headers249 250    @property251    def headers(self) -> HTTPHeadersDict:252        return self._headers.copy()253 254    @property255    def cookies(self) -> RequestsCookieJar:256        self.cookie_jar.clear_expired_cookies()257        return self.cookie_jar258 259    @cookies.setter260    def cookies(self, jar: RequestsCookieJar):261        self.cookie_jar = jar262 263    def remove_cookies(self, cookies: List[Dict[str, str]]):264        for cookie in cookies:265            remove_cookie_by_name(266                self.cookie_jar,267                cookie['name'],268                domain=cookie.get('domain', None),269                path=cookie.get('path', None)270            )271 272    @property273    def auth(self) -> Optional[AuthBase]:274        auth = self.get('auth', None)275        if not auth or not auth['type']:276            return277 278        plugin = plugin_manager.get_auth_plugin(auth['type'])()279 280        credentials = {'username': None, 'password': None}281        try:282            # New style283            plugin.raw_auth = auth['raw_auth']284        except KeyError:285            # Old style286            credentials = {287                'username': auth['username'],288                'password': auth['password'],289            }290        else:291            if plugin.auth_parse:292                from .cli.argtypes import parse_auth293                parsed = parse_auth(plugin.raw_auth)294                credentials = {295                    'username': parsed.key,296                    'password': parsed.value,297                }298 299        return plugin.get_auth(**credentials)300 301    @auth.setter302    def auth(self, auth: dict):303        assert {'type', 'raw_auth'} == auth.keys()304        self['auth'] = auth305 306    @property307    def is_anonymous(self):308        return is_anonymous_session(self.session_id)309 310    def warn_legacy_usage(self, warning: str) -> None:311        if self.suppress_legacy_warnings:312            return None313 314        self.env.log_error(315            warning,316            level=LogLevel.WARNING317        )318 319        # We don't want to spam multiple warnings on each usage,320        # so if there is already a warning for the legacy usage321        # we'll skip the next ones.322        self.suppress_legacy_warnings = True323