"""Persistent, JSON-serialized sessions. """ import os import sys import json import glob import errno import codecs import shutil import subprocess from requests.compat import urlparse from requests import Session as RSession from requests.cookies import RequestsCookieJar, create_cookie from requests.auth import HTTPBasicAuth, HTTPDigestAuth 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): 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 = RSession(cookies=session.cookies) try: response = rsession.request(**request_kwargs) except Exception: raise else: session.cookies = rsession.cookies session.save() return response class Host(object): def __init__(self, name): self.name = name def __iter__(self): 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): 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'] = {} @property def path(self): return os.path.join(self.host.path, self.name + '.json') 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 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='?') # Show show = subparsers.add_parser('session-show', help='list or show sessions') 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='?', help='The name of the session to be deleted.' ' If not specified, all host sessions are deleted.')