httpie-cli/tests/test_sessions.py

# coding=utf-8
import json
import os
import shutil
from datetime import datetime
from mock import mock
from tempfile import gettempdir

import pytest

from fixtures import UNICODE
from httpie.plugins import AuthPlugin
from httpie.plugins.builtin import HTTPBasicAuth
from httpie.plugins.registry import plugin_manager
from httpie.sessions import Session
from httpie.utils import get_expired_cookies
from tests.test_auth_plugins import basic_auth
from utils import HTTP_OK, MockEnvironment, http, mk_config_dir


class SessionTestBase:

    def start_session(self, httpbin):
        """Create and reuse a unique config dir for each test."""
        self.config_dir = mk_config_dir()

    def teardown_method(self, method):
        shutil.rmtree(self.config_dir)

    def env(self):
        """
        Return an environment.

        Each environment created within a test method
        will share the same config_dir. It is necessary
        for session files being reused.

        """
        return MockEnvironment(config_dir=self.config_dir)


class CookieTestBase:
    def setup_method(self, method):
        self.config_dir = mk_config_dir()

        orig_session = {
            'cookies': {
                'cookie1': {
                    'value': 'foo',
                },
                'cookie2': {
                    'value': 'foo',
                }
            }
        }
        self.session_path = self.config_dir / 'test-session.json'
        self.session_path.write_text(json.dumps(orig_session))

    def teardown_method(self, method):
        shutil.rmtree(self.config_dir)


class TestSessionFlow(SessionTestBase):
    """
    These tests start with an existing session created in `setup_method()`.

    """

    def start_session(self, httpbin):
        """
        Start a full-blown session with a custom request header,
        authorization, and response cookies.

        """
        super().start_session(httpbin)
        r1 = http(
            '--follow',
            '--session=test',
            '--auth=username:password',
            'GET',
            httpbin.url + '/cookies/set?hello=world',
            'Hello:World',
            env=self.env()
        )
        assert HTTP_OK in r1

    def test_session_created_and_reused(self, httpbin):
        self.start_session(httpbin)
        # Verify that the session created in setup_method() has been used.
        r2 = http('--session=test',
                  'GET', httpbin.url + '/get', env=self.env())
        assert HTTP_OK in r2
        assert r2.json['headers']['Hello'] == 'World'
        assert r2.json['headers']['Cookie'] == 'hello=world'
        assert 'Basic ' in r2.json['headers']['Authorization']

    def test_session_update(self, httpbin):
        self.start_session(httpbin)
        # Get a response to a request from the original session.
        r2 = http('--session=test', 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r2

        # Make a request modifying the session data.
        r3 = http('--follow', '--session=test', '--auth=username:password2',
                  'GET', httpbin.url + '/cookies/set?hello=world2',
                  'Hello:World2',
                  env=self.env())
        assert HTTP_OK in r3

        # Get a response to a request from the updated session.
        r4 = http('--session=test', 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r4
        assert r4.json['headers']['Hello'] == 'World2'
        assert r4.json['headers']['Cookie'] == 'hello=world2'
        assert (r2.json['headers']['Authorization']
                != r4.json['headers']['Authorization'])

    def test_session_read_only(self, httpbin):
        self.start_session(httpbin)
        # Get a response from the original session.
        r2 = http('--session=test', 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r2

        # Make a request modifying the session data but
        # with --session-read-only.
        r3 = http('--follow', '--session-read-only=test',
                  '--auth=username:password2', 'GET',
                  httpbin.url + '/cookies/set?hello=world2', 'Hello:World2',
                  env=self.env())
        assert HTTP_OK in r3

        # Get a response from the updated session.
        r4 = http('--session=test', 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r4

        # Origin can differ on Travis.
        del r2.json['origin'], r4.json['origin']
        # Different for each request.

        # Should be the same as before r3.
        assert r2.json == r4.json


class TestSession(SessionTestBase):
    """Stand-alone session tests."""

    def test_session_ignored_header_prefixes(self, httpbin):
        self.start_session(httpbin)
        r1 = http('--session=test', 'GET', httpbin.url + '/get',
                  'Content-Type: text/plain',
                  'If-Unmodified-Since: Sat, 29 Oct 1994 19:43:31 GMT',
                  env=self.env())
        assert HTTP_OK in r1
        r2 = http('--session=test', 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r2
        assert 'Content-Type' not in r2.json['headers']
        assert 'If-Unmodified-Since' not in r2.json['headers']

    def test_session_by_path(self, httpbin):
        self.start_session(httpbin)
        session_path = self.config_dir / 'session-by-path.json'
        r1 = http('--session', str(session_path), 'GET', httpbin.url + '/get',
                  'Foo:Bar', env=self.env())
        assert HTTP_OK in r1

        r2 = http('--session', str(session_path), 'GET', httpbin.url + '/get',
                  env=self.env())
        assert HTTP_OK in r2
        assert r2.json['headers']['Foo'] == 'Bar'

    def test_session_unicode(self, httpbin):
        self.start_session(httpbin)

        r1 = http('--session=test', u'--auth=test:' + UNICODE,
                  'GET', httpbin.url + '/get', u'Test:%s' % UNICODE,
                  env=self.env())
        assert HTTP_OK in r1

        r2 = http('--session=test', '--verbose', 'GET',
                  httpbin.url + '/get', env=self.env())
        assert HTTP_OK in r2

        # FIXME: Authorization *sometimes* is not present on Python3
        assert (r2.json['headers']['Authorization']
                == HTTPBasicAuth.make_header(u'test', UNICODE))
        # httpbin doesn't interpret utf8 headers
        assert UNICODE in r2

    def test_session_default_header_value_overwritten(self, httpbin):
        self.start_session(httpbin)
        # https://github.com/jakubroztocil/httpie/issues/180
        r1 = http('--session=test',
                  httpbin.url + '/headers', 'User-Agent:custom',
                  env=self.env())
        assert HTTP_OK in r1
        assert r1.json['headers']['User-Agent'] == 'custom'

        r2 = http('--session=test', httpbin.url + '/headers', env=self.env())
        assert HTTP_OK in r2
        assert r2.json['headers']['User-Agent'] == 'custom'

    def test_download_in_session(self, httpbin):
        # https://github.com/jakubroztocil/httpie/issues/412
        self.start_session(httpbin)
        cwd = os.getcwd()
        os.chdir(gettempdir())
        try:
            http('--session=test', '--download',
                 httpbin.url + '/get', env=self.env())
        finally:
            os.chdir(cwd)

    @pytest.mark.parametrize(
        argnames=['auth_require_param', 'auth_parse_param'],
        argvalues=[
            (False, False),
            (False, True),
            (True, False)
        ]
    )
    def test_auth_type_reused_in_session(self, auth_require_param, auth_parse_param, httpbin):
        self.start_session(httpbin)
        session_path = self.config_dir / 'test-session.json'

        header = 'Custom dXNlcjpwYXNzd29yZA'

        class Plugin(AuthPlugin):
            auth_type = 'test-reused'
            auth_require = auth_require_param
            auth_parse = auth_parse_param

            def get_auth(self, username=None, password=None):
                return basic_auth(header=f'{header}==')

        plugin_manager.register(Plugin)

        r1 = http(
            '--session', str(session_path),
            httpbin + '/basic-auth/user/password',
            '--auth-type',
            Plugin.auth_type,
            '--auth', 'user:password',
            '--print=H',
        )

        r2 = http(
            '--session', str(session_path),
            httpbin + '/basic-auth/user/password',
            '--print=H',
        )
        assert f'Authorization: {header}' in r1
        assert f'Authorization: {header}' in r2
        plugin_manager.unregister(Plugin)

    @mock.patch('httpie.cli.argtypes.AuthCredentials._getpass',
                new=lambda self, prompt: 'password')
    def test_auth_plugin_prompt_password_in_session(self, httpbin):
        self.start_session(httpbin)
        session_path = self.config_dir / 'test-session.json'

        class Plugin(AuthPlugin):
            auth_type = 'test-prompted'

            def get_auth(self, username=None, password=None):
                return basic_auth()

        plugin_manager.register(Plugin)

        r1 = http(
            '--session', str(session_path),
            httpbin + '/basic-auth/user/password',
            '--auth-type',
            Plugin.auth_type,
            '--auth', 'user:',
        )

        r2 = http(
            '--session', str(session_path),
            httpbin + '/basic-auth/user/password',
        )
        assert HTTP_OK in r1
        assert HTTP_OK in r2
        plugin_manager.unregister(Plugin)

    def test_auth_type_stored_in_session_file(self, httpbin):
        self.config_dir = mk_config_dir()
        self.session_path = self.config_dir / 'test-session.json'

        class Plugin(AuthPlugin):
            auth_type = 'test-saved'
            auth_require = True

            def get_auth(self, username=None, password=None):
                return basic_auth()

        plugin_manager.register(Plugin)
        http('--session', str(self.session_path),
             httpbin + '/basic-auth/user/password',
             '--auth-type',
             Plugin.auth_type,
             '--auth', 'user:password',
             )
        updated_session = json.loads(self.session_path.read_text())
        assert updated_session['auth']['type'] == 'test-saved'
        assert updated_session['auth']['raw_auth'] == "user:password"
        plugin_manager.unregister(Plugin)


class TestExpiredCookies(CookieTestBase):

    @pytest.mark.parametrize(
        argnames=['initial_cookie', 'expired_cookie'],
        argvalues=[
            ({'id': {'value': 123}}, 'id'),
            ({'id': {'value': 123}}, 'token')
        ]
    )
    def test_removes_expired_cookies_from_session_obj(self, initial_cookie, expired_cookie, httpbin):
        session = Session(self.config_dir)
        session['cookies'] = initial_cookie
        session.remove_cookies([expired_cookie])
        assert expired_cookie not in session.cookies

    def test_expired_cookies(self, httpbin):
        r = http(
            '--session', str(self.session_path),
            '--print=H',
            httpbin.url + '/cookies/delete?cookie2',
        )
        assert 'Cookie: cookie1=foo; cookie2=foo' in r

        updated_session = json.loads(self.session_path.read_text())
        assert 'cookie1' in updated_session['cookies']
        assert 'cookie2' not in updated_session['cookies']

    @pytest.mark.parametrize(
        argnames=['headers', 'now', 'expected_expired'],
        argvalues=[
            (
                [
                    ('Set-Cookie', 'hello=world; Path=/; Expires=Thu, 01-Jan-1970 00:00:00 GMT; HttpOnly'),
                    ('Connection', 'keep-alive')
                ],
                None,
                [
                    {
                        'name': 'hello',
                        'path': '/'
                    }
                ]
            ),
            (
                [
                    ('Set-Cookie', 'hello=world; Path=/; Expires=Thu, 01-Jan-1970 00:00:00 GMT; HttpOnly'),
                    ('Set-Cookie', 'pea=pod; Path=/ab; Expires=Thu, 01-Jan-1970 00:00:00 GMT; HttpOnly'),
                    ('Connection', 'keep-alive')
                ],
                None,
                [
                    {'name': 'hello', 'path': '/'},
                    {'name': 'pea', 'path': '/ab'}
                ]
            ),
            (
                [
                    ('Set-Cookie', 'hello=world; Path=/; Expires=Fri, 12 Jun 2020 12:28:55 GMT; HttpOnly'),
                    ('Connection', 'keep-alive')
                ],
                datetime(2020, 6, 11).timestamp(),
                []
            )
        ]
    )
    def test_get_expired_cookies_manages_multiple_cookie_headers(self, headers, now, expected_expired):
        assert get_expired_cookies(headers, now=now) == expected_expired


class TestCookieStorage(CookieTestBase):

    @pytest.mark.parametrize(
        argnames=['new_cookies', 'new_cookies_dict', 'expected'],
        argvalues=[(
            'new=bar',
            {'new': 'bar'},
            'cookie1=foo; cookie2=foo; new=bar'
        ),
            (
            'new=bar;chocolate=milk',
            {'new': 'bar', 'chocolate': 'milk'},
            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
        ),
            (
            'new=bar; chocolate=milk',
            {'new': 'bar', 'chocolate': 'milk'},
            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
        ),
            (
            'new=bar;; chocolate=milk;;;',
            {'new': 'bar', 'chocolate': 'milk'},
            'cookie1=foo; cookie2=foo; new=bar'
        ),
            (
            'new=bar; chocolate=milk;;;',
            {'new': 'bar', 'chocolate': 'milk'},
            'chocolate=milk; cookie1=foo; cookie2=foo; new=bar'
        )
        ]
    )
    def test_existing_and_new_cookies_sent_in_request(self, new_cookies, new_cookies_dict, expected, httpbin):
        r = http(
            '--session', str(self.session_path),
            '--print=H',
            httpbin.url,
            'Cookie:' + new_cookies,
        )
        # Note: cookies in response are in alphabetical order
        assert 'Cookie: ' + expected in r

        updated_session = json.loads(self.session_path.read_text())
        for name, value in new_cookies_dict.items():
            assert name, value in updated_session['cookies']
            assert 'Cookie' not in updated_session['headers']

    @pytest.mark.parametrize(
        argnames=['cli_cookie', 'set_cookie', 'expected'],
        argvalues=[(
            '',
            '/cookies/set/cookie1/bar',
            'bar'
        ),
            (
            'cookie1=not_foo',
            '/cookies/set/cookie1/bar',
            'bar'
        ),
            (
            'cookie1=not_foo',
            '',
            'not_foo'
        ),
            (
            '',
            '',
            'foo'
        )
        ]
    )
    def test_cookie_storage_priority(self, cli_cookie, set_cookie, expected, httpbin):
        """
        Expected order of priority for cookie storage in session file:
        1. set-cookie (from server)
        2. command line arg
        3. cookie already stored in session file
        """
        r = http(
            '--session', str(self.session_path),
            httpbin.url + set_cookie,
            'Cookie:' + cli_cookie,
        )
        updated_session = json.loads(self.session_path.read_text())

        assert updated_session['cookies']['cookie1']['value'] == expected