httpie-cli/httpie/sessions.py

"""Persistent, JSON-serialized sessions.

"""
import os
import sys
import json
import glob
import errno
import codecs
import shutil
import subprocess

import requests
from requests.compat import urlparse
from requests.cookies import RequestsCookieJar, create_cookie
from requests.auth import HTTPBasicAuth, HTTPDigestAuth
from argparse import OPTIONAL

from.import __version__
from.config import CONFIG_DIR
from.output import PygmentsProcessor


SESSIONS_DIR = os.path.join(CONFIG_DIR, 'sessions')


def get_response(name, request_kwargs, read_only=False):
    """Like `client.get_response`, but applies permanent
    aspects of the session to the request.

    """
    host = Host(request_kwargs['headers'].get('Host', None)
    or urlparse(request_kwargs['url']).netloc.split('@')[-1])

    session = Session(host, name)
    session.load()

    # Update session headers with the request headers.
    session['headers'].update(request_kwargs.get('headers', {}))
    # Use the merged headers for the request
    request_kwargs['headers'] = session['headers']

    auth = request_kwargs.get('auth', None)
    if auth:
        session.auth = auth
    elif session.auth:
        request_kwargs['auth'] = session.auth

    rsession = requests.Session(cookies=session.cookies)
    try:
        response = rsession.request(**request_kwargs)
    except Exception:
        raise
    else:
        if not read_only or session.is_new:
            session.cookies = rsession.cookies
            session.save()
        return response


class Host(object):
    """A host is a per-host directory on the disk containing sessions files."""

    def __init__(self, name):
        self.name = name

    def __iter__(self):
        """Return a iterator yielding `(session_name, session_path)`."""
        for fn in sorted(glob.glob1(self.path, '*.json')):
            yield os.path.splitext(fn)[0], os.path.join(self.path, fn)

    def delete(self):
        shutil.rmtree(self.path)

    @property
    def path(self):
        # Name will include ':' if a port is specified, which is invalid
        # on windows. DNS does not allow '_' in a domain, or for it to end
        # in a number (I think?)
        path = os.path.join(SESSIONS_DIR, self.name.replace(':', '_'))
        try:
            os.makedirs(path, mode=0o700)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
        return path

    @classmethod
    def all(cls):
        """Return a generator yielding a host at a time."""
        for name in sorted(glob.glob1(SESSIONS_DIR, '*')):
            if os.path.isdir(os.path.join(SESSIONS_DIR, name)):
                yield Host(name)


class Session(dict):
    """"""

    def __init__(self, host, name, *args, **kwargs):
        super(Session, self).__init__(*args, **kwargs)
        self.host = host
        self.name = name
        self['headers'] = {}
        self['cookies'] = {}

    def load(self):
        try:
            with open(self.path, 'rt') as f:
                try:
                    data = json.load(f)
                except ValueError as e:
                    raise ValueError('Invalid session: %s [%s]' %
                                     (e.message, self.path))
                self.update(data)
        except IOError as e:
            if e.errno != errno.ENOENT:
                raise

    def save(self):
        self['__version__'] = __version__
        with open(self.path, 'w') as f:
            json.dump(self, f, indent=4, sort_keys=True, ensure_ascii=True)
            f.write('\n')

    def delete(self):
        try:
            os.unlink(self.path)
        except OSError as e:
            if e.errno != errno.ENOENT:
                raise

    @property
    def path(self):
        return os.path.join(self.host.path, self.name + '.json')

    @property
    def is_new(self):
        return not os.path.exists(self.path)

    @property
    def cookies(self):
        jar = RequestsCookieJar()
        for name, cookie_dict in self['cookies'].items():
            jar.set_cookie(create_cookie(
                name, cookie_dict.pop('value'), **cookie_dict))
        jar.clear_expired_cookies()
        return jar

    @cookies.setter
    def cookies(self, jar):
        excluded = [
            '_rest', 'name', 'port_specified',
            'domain_specified', 'domain_initial_dot',
            'path_specified', 'comment', 'comment_url'
        ]
        self['cookies'] = {}
        for host in jar._cookies.values():
            for path in host.values():
                for name, cookie in path.items():
                    cookie_dict = {}
                    for k, v in cookie.__dict__.items():
                        if k not in excluded:
                            cookie_dict[k] = v
                    self['cookies'][name] = cookie_dict

    @property
    def auth(self):
        auth = self.get('auth', None)
        if not auth:
            return None
        Auth = {'basic': HTTPBasicAuth,
                'digest': HTTPDigestAuth}[auth['type']]
        return Auth(auth['username'], auth['password'])

    @auth.setter
    def auth(self, cred):
        self['auth'] = {
            'type': {HTTPBasicAuth: 'basic',
                     HTTPDigestAuth: 'digest'}[type(cred)],
            'username': cred.username,
            'password': cred.password,
        }


def list_command(args):
    if args.host:
        for name, path in Host(args.host):
            print(name + ' [' + path + ']')
    else:
        for host in Host.all():
            print(host.name)
            for name, path in host:
                print(' ' + name + ' [' + path + ']')


def show_command(args):
    path = Session(Host(args.host), args.name).path
    if not os.path.exists(path):
        sys.stderr.write('Session "%s" does not exist [%s].\n'
        % (args.name, path))
        sys.exit(1)

    with codecs.open(path, encoding='utf8') as f:
        print(path + ':\n')
        proc = PygmentsProcessor()
        print(proc.process_body(f.read(), 'application/json', 'json'))
        print('')


def delete_command(args):
    host = Host(args.host)
    if not args.name:
        host.delete()
    else:
        Session(host, args.name).delete()


def edit_command(args):
    editor = os.environ.get('EDITOR', None)
    if not editor:
        sys.stderr.write(
            'You need to configure the environment variable EDITOR.\n')
        sys.exit(1)
    command = editor.split()
    command.append(Session(Host(args.host), args.name).path)
    subprocess.call(command)


def add_commands(subparsers):

    # List
    list_ = subparsers.add_parser('session-list', help='list sessions')
    list_.set_defaults(command=list_command)
    list_.add_argument('host', nargs=OPTIONAL)

    # Show
    show = subparsers.add_parser('session-show', help='show a session')
    show.set_defaults(command=show_command)
    show.add_argument('host')
    show.add_argument('name')

    # Edit
    edit = subparsers.add_parser(
        'session-edit', help='edit a session in $EDITOR')
    edit.set_defaults(command=edit_command)
    edit.add_argument('host')
    edit.add_argument('name')

    # Delete
    delete = subparsers.add_parser('session-delete', help='delete a session')
    delete.set_defaults(command=delete_command)
    delete.add_argument('host')
    delete.add_argument('name', nargs=OPTIONAL,
        help='The name of the session to be deleted.'
             ' If not specified, all host sessions are deleted.')