<commit_before># Copyright (C) 2017 Bradley Sakdol <bsakdol@turnitin.com>
#
# This file is part of Jumpbox
#
# Jumpbox is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""This module handles the Jumpbox UI"""

import curses
import curses.panel
import os

from menu import Menu
from _version import __version__

class JumpboxUI(object):
    def __init__(self, stdscr):
        self.stdscr = stdscr

        curses.start_color()
        curses.noecho()
        curses.cbreak()
        curses.curs_set(0)
        self.stdscr.keypad(1)

        curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_WHITE)
        self.highlight = curses.color_pair(1)
        self.normal = curses.A_NORMAL

        self.processmenu(Menu().main_menu)

    def processmenu(self, menu, parent=None):
        optioncount = len(menu['options'])
        exitmenu = False

        while not exitmenu:
            selected = self.runmenu(menu, parent)

            if selected == optioncount:
                exitmenu = True

            elif menu['options'][selected]['type'] == "menu":
                self.stdscr.clear()
                self.processmenu(menu['options'][selected], menu)
                self.stdscr.clear()

            elif menu['options'][selected]['type'] == "command":
                curses.def_prog_mode()
                os.system('reset')
                username = raw_input('Username: ')
                os.system('ssh ' + username + '@' + menu['options'][selected][
                    'ip_addr'])
                self.stdscr.clear()
                curses.reset_prog_mode()
                curses.curs_set(1)
                curses.curs_set(0)

            elif menu['options'][selected]['type'] == "prompt":
                curses.def_prog_mode()
                os.system('reset')
                ip_address = raw_input('IP / Hostname: ')
                os.system('clear')
                username = raw_input('Username: ')
                os.system('ssh ' + username + '@' + ip_address)
                self.stdscr.clear()
                curses.reset_prog_mode()
                curses.curs_set(1)
                curses.curs_set(0)

    def runmenu(self, menu, parent):
        self.stdscr.refresh()

        if parent is None:
            lastoption = 'Exit'
        else:
            lastoption = 'Return to the %s menu' % parent['title']

        optioncount = len(menu['options'])
        pos = 0
        oldpos = None
        scroller = 0
        cmd = None

        termy, termx = self.stdscr.getmaxyx()
        viewport = termy - 8
        height, width = optioncount + 3, termx - 4
        top, left = 5, 3

        pad = curses.newpad(height, width)

        while cmd != ord('\n'):
            if pos != oldpos:
                oldpos = pos
                self.stdscr.border('|', '|', '-', '-', '+', '+', '+', '+')
                self.stdscr.addstr(1, 2, 'NetOps Jumpbox', curses.A_UNDERLINE)
                self.stdscr.addstr(2, 3, menu['title'], curses.A_BOLD)
                self.stdscr.addstr(4, 3, menu['subtitle'], curses.A_NORMAL)
                self.stdscr.addstr(termy - 2, termx - 9, 'v' + __version__,
                                   curses.A_BOLD)
                self.stdscr.refresh()

                for index in range(optioncount):
                    textstyle = self.normal
                    if pos == index:
                        textstyle = self.highlight
                    if 'ip_addr' in menu['options'][index]:
                        pad.addstr(index, 0, '%d - %s' % (
                            index + 1, menu['options'][index]['title'] + ': ' +
                            menu['options'][index]['ip_addr']), textstyle)
                    else:
                        pad.addstr(index, 0, '%d - %s' % (
                            index + 1, menu['options'][index]['title']),
                            textstyle)

                textstyle = self.normal
                if pos == optioncount:
                    textstyle = self.highlight
                pad.addstr(optioncount, 0, '%d - %s' %
                           (optioncount + 1, lastoption), textstyle)

                pad.refresh(scroller, 0, top, left, termy - 3, width)

            cmd = self.stdscr.getch()

            if cmd == curses.KEY_RESIZE:
                self.stdscr.erase()
                termy, termx = self.stdscr.getmaxyx()
                viewport = termy - 8
                width = termx - 4
                self.stdscr.border('|', '|', '-', '-', '+', '+', '+', '+')
                self.stdscr.addstr(1, 2, 'NetOps Jumpbox', curses.A_UNDERLINE)
                self.stdscr.addstr(2, 3, menu['title'], curses.A_BOLD)
                self.stdscr.addstr(4, 3, menu['subtitle'], curses.A_NORMAL)
                self.stdscr.addstr(termy - 2, termx - 9, 'v' + __version__,
                                   curses.A_BOLD)
                self.stdscr.refresh()
                pad.refresh(scroller, 0, top, left, termy - 3, width)

            if optioncount + 1 > viewport:
                if cmd == curses.KEY_DOWN:
                    if pos < optioncount and scroller < height - viewport:
                        pos += 1
                        scroller += 1
                    elif pos < optioncount and scroller == height - viewport:
                        pos += 1
                    elif pos == optioncount:
                        pos = 0
                        scroller = 0

                elif cmd == curses.KEY_UP:
                    if pos > 0 and scroller > 0:
                        pos -= 1
                        scroller -= 1
                    elif pos > 0 and scroller == 0:
                        pos -= 1
                    elif pos == 0:
                        pos = optioncount
                        scroller = height - viewport

                elif cmd == curses.KEY_RIGHT:
                    cmd = ord('\n')

                elif cmd == curses.KEY_LEFT:
                    pos = optioncount
                    cmd = ord('\n')

            else:
                if cmd == curses.KEY_DOWN:
                    if pos < optioncount:
                        pos += 1
                    else:
                        pos = 0

                elif cmd == curses.KEY_UP:
                    if pos > 0:
                        pos -= 1
                    else:
                        pos = optioncount

                elif cmd == curses.KEY_RIGHT:
                    cmd = ord('\n')

                elif cmd == curses.KEY_LEFT:
                    pos = optioncount
                    cmd = ord('\n')

        return pos
<commit_msg>format for pep8<commit_after># Copyright (C) 2017 Bradley Sakdol <bsakdol@turnitin.com>
#
# This file is part of Jumpbox
#
# Jumpbox is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""This module handles the Jumpbox UI"""

import curses
import curses.panel
import os

from menu import Menu
from _version import __version__

class JumpboxUI(object):
    def __init__(self, stdscr):
        self.stdscr = stdscr

        curses.start_color()
        curses.noecho()
        curses.cbreak()
        curses.curs_set(0)
        self.stdscr.keypad(1)

        curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_WHITE)
        self.highlight = curses.color_pair(1)
        self.normal = curses.A_NORMAL

        self.processmenu(Menu().main_menu)

    def processmenu(self, menu, parent=None):
        optioncount = len(menu['options'])
        exitmenu = False

        while not exitmenu:
            selected = self.runmenu(menu, parent)

            if selected == optioncount:
                exitmenu = True

            elif menu['options'][selected]['type'] == "menu":
                self.stdscr.clear()
                self.processmenu(menu['options'][selected], menu)
                self.stdscr.clear()

            elif menu['options'][selected]['type'] == "command":
                curses.def_prog_mode()
                os.system('reset')
                username = raw_input('Username: ')
                os.system('ssh ' + username + '@' +
                          menu['options'][selected]['ip_addr'])
                self.stdscr.clear()
                curses.reset_prog_mode()
                curses.curs_set(1)
                curses.curs_set(0)

            elif menu['options'][selected]['type'] == "prompt":
                curses.def_prog_mode()
                os.system('reset')
                ip_address = raw_input('IP / Hostname: ')
                os.system('clear')
                username = raw_input('Username: ')
                os.system('ssh ' + username + '@' + ip_address)
                self.stdscr.clear()
                curses.reset_prog_mode()
                curses.curs_set(1)
                curses.curs_set(0)

    def runmenu(self, menu, parent):
        self.stdscr.refresh()

        if parent is None:
            lastoption = 'Exit'
        else:
            lastoption = 'Return to the %s menu' % parent['title']

        optioncount = len(menu['options'])
        pos = 0
        oldpos = None
        scroller = 0
        cmd = None

        termy, termx = self.stdscr.getmaxyx()
        viewport = termy - 8
        height, width = optioncount + 3, termx - 4
        top, left = 5, 3

        pad = curses.newpad(height, width)

        while cmd != ord('\n'):
            if pos != oldpos:
                oldpos = pos
                self.stdscr.border('|', '|', '-', '-', '+', '+', '+', '+')
                self.stdscr.addstr(1, 2, 'NetOps Jumpbox', curses.A_UNDERLINE)
                self.stdscr.addstr(2, 3, menu['title'], curses.A_BOLD)
                self.stdscr.addstr(4, 3, menu['subtitle'], curses.A_NORMAL)
                self.stdscr.addstr(termy - 2, termx - 9, 'v' + __version__,
                                   curses.A_BOLD)
                self.stdscr.refresh()

                for index in range(optioncount):
                    textstyle = self.normal
                    if pos == index:
                        textstyle = self.highlight
                    if 'ip_addr' in menu['options'][index]:
                        pad.addstr(
                            index, 0, '%d - %s' %
                            (index + 1, menu['options'][index]['title'] + ': '
                             + menu['options'][index]['ip_addr']), textstyle)
                    else:
                        pad.addstr(index, 0, '%d - %s' %
                                   (index + 1,
                                    menu['options'][index]['title']),
                                   textstyle)

                textstyle = self.normal
                if pos == optioncount:
                    textstyle = self.highlight
                pad.addstr(optioncount, 0, '%d - %s' % (optioncount + 1,
                                                        lastoption), textstyle)

                pad.refresh(scroller, 0, top, left, termy - 3, width)

            cmd = self.stdscr.getch()

            if cmd == curses.KEY_RESIZE:
                self.stdscr.erase()
                termy, termx = self.stdscr.getmaxyx()
                viewport = termy - 8
                width = termx - 4
                self.stdscr.border('|', '|', '-', '-', '+', '+', '+', '+')
                self.stdscr.addstr(1, 2, 'NetOps Jumpbox', curses.A_UNDERLINE)
                self.stdscr.addstr(2, 3, menu['title'], curses.A_BOLD)
                self.stdscr.addstr(4, 3, menu['subtitle'], curses.A_NORMAL)
                self.stdscr.addstr(termy - 2, termx - 9, 'v' + __version__,
                                   curses.A_BOLD)
                self.stdscr.refresh()
                pad.refresh(scroller, 0, top, left, termy - 3, width)

            if optioncount + 1 > viewport:
                if cmd == curses.KEY_DOWN:
                    if pos < optioncount and scroller < height - viewport:
                        pos += 1
                        scroller += 1
                    elif pos < optioncount and scroller == height - viewport:
                        pos += 1
                    elif pos == optioncount:
                        pos = 0
                        scroller = 0

                elif cmd == curses.KEY_UP:
                    if pos > 0 and scroller > 0:
                        pos -= 1
                        scroller -= 1
                    elif pos > 0 and scroller == 0:
                        pos -= 1
                    elif pos == 0:
                        pos = optioncount
                        scroller = height - viewport

                elif cmd == curses.KEY_RIGHT:
                    cmd = ord('\n')

                elif cmd == curses.KEY_LEFT:
                    pos = optioncount
                    cmd = ord('\n')

            else:
                if cmd == curses.KEY_DOWN:
                    if pos < optioncount:
                        pos += 1
                    else:
                        pos = 0

                elif cmd == curses.KEY_UP:
                    if pos > 0:
                        pos -= 1
                    else:
                        pos = optioncount

                elif cmd == curses.KEY_RIGHT:
                    cmd = ord('\n')

                elif cmd == curses.KEY_LEFT:
                    pos = optioncount
                    cmd = ord('\n')

        return pos
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import subprocess
import sys

import boto3
import botocore
import click

def format_output(instances, flag):
    """return formatted string for instance"""
    out = []
    line_format = '{0}\t{1}\t{2}\t{3}\t{4}'
    name_len = _get_max_name_len(instances) + 3
    if flag:
        line_format = '{0:<' + str(name_len) + '}{1:<10}{2:<13}{3:<16}{4:<16}'

    for i in instances:
        tag_name = get_tag_value(i.tags, 'Name')
        out.append(line_format.format(
            tag_name, i.state['Name'], i.id, i.private_ip_address, str(i.public_ip_address)))
    return out

def _get_max_name_len(instances):
    """get max length of Tag:Name"""
    # FIXME: ec2.instanceCollection doesn't have __len__
    for i in instances:
        return max([len(get_tag_value(i.tags, 'Name')) for i in instances])
    return 0

def get_tag_value(x, key):
    """Get a value from tag"""
    if x is None:
        return ''
    result = [y['Value'] for y in x if y['Key'] == key]
    if result:
        return result[0]
    return ''

@click.group()
def cli():
    """EC2 CLI group"""
    pass

@cli.command(help='List EC2 instances')
@click.argument('name', default='*')
@click.option('--list-formatted', '-l', is_flag=True)
def ls(name, list_formatted):
    """List EC2 instances"""
    ec2 = boto3.resource('ec2')
    if name == '*':
        instances = ec2.instances.filter()
    else:
        condition = {'Name': 'tag:Name', 'Values': [name]}
        instances = ec2.instances.filter(Filters=[condition])
    out = format_output(instances, list_formatted)
    click.echo('\n'.join(out))

@cli.command(help='Start EC2 instance')
@click.option('--instance-id', '-i', required=True, help='EC2 instance id')
def up(instance_id):
    """Start EC2 instance"""
    ec2 = boto3.resource('ec2')
    try:
        instance = ec2.Instance(instance_id)
        instance.start()
    except botocore.exceptions.ClientError as e:
        click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
        sys.exit(2)

@cli.command(help='Stop EC2 instance')
@click.option('--instance-id', '-i', required=True, help='EC2 instance id')
def down(instance_id):
    """Stop EC2 instance"""
    ec2 = boto3.resource('ec2')
    try:
        instance = ec2.Instance(instance_id)
        instance.stop()
    except botocore.exceptions.ClientError as e:
        click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
        sys.exit(2)

def create_ssh_command(instance_id, instance_name, username, key_file, port, gateway_instance_id):
    """Create SSH Login command string"""
    ec2 = boto3.resource('ec2')
    if instance_id is not None:
        try:
            instance = ec2.Instance(instance_id)
            hostname = instance.public_ip_address
        except botocore.exceptions.ClientError as e:
            click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
            sys.exit(2)
    elif instance_name is not None:
        try:
            condition = {'Name': 'tag:Name', 'Values': [instance_name]}
            instances = ec2.instances.filter(Filters=[condition])
            target_instances = []
            for idx, i in enumerate(instances):
                target_instances.append(i)
                tag_name = get_tag_value(i.tags, 'Name')
                click.echo('[{0}]: {1}\t{2}\t{3}\t{4}'.format(
                    idx, i.id, i.private_ip_address, i.state['Name'], tag_name))
            selected_idx = click.prompt("Please enter a valid number", type=int, default=0)
            # TODO: add validation for if selected_idx exceeds length of target_instances
            click.echo("{0} is selected.".format(selected_idx))
            instance = target_instances[selected_idx]
            hostname = instance.public_ip_address
        except botocore.exceptions.ClientError as e:
            click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
            sys.exit(2)
    # TODO: need to refactor and make it testable
    if gateway_instance_id is not None:
        gateway_instance = ec2.Instance(gateway_instance_id)
        gateway_public_ip = gateway_instance.public_ip_address
        hostname = instance.private_ip_address
        cmd = 'ssh -tt ubuntu@{0} -i {1} -p {2} ssh {3}@{4}'.format(
            gateway_public_ip, key_file, port, username, hostname)
    else:
        cmd = 'ssh {0}@{1} -i {2} -p {3}'.format(username, hostname, key_file, port)
    return cmd

@cli.command(help='SSH login to EC2 instance')
@click.option('--instance-id', '-i', default=None, help='EC2 instance id')
@click.option('--instance-name', '-n', default=None, help='EC2 instance Name Tag')
@click.option('--username', '-u', default='ubuntu', help='Login username')
@click.option('--key-file', '-k', required=True, help='SSH Key file path', type=click.Path())
@click.option('--port', '-p', help='SSH port', default=22)
@click.option('--gateway-instance-id', '-g', default=None, help='Gateway instance id')
@click.option('--dry-run', is_flag=True, default=False, help='Print SSH Login command and exist')
def ssh(instance_id, instance_name, username, key_file, port, gateway_instance_id, dry_run):
    """SSH to EC2 instance"""
    if instance_id is None and instance_name is None:
        click.echo(
            "One of --instance-id/-i or --instance-name/-n"
            " has to be specified.", err=True)
        sys.exit(2)
    elif instance_id is not None and instance_name is not None:
        click.echo(
            "Both --instance-id/-i and --instance-name/-n "
            "can't to be specified at the same time.", err=True)
        sys.exit(2)
    cmd = create_ssh_command(
        instance_id, instance_name, username, key_file, port, gateway_instance_id)
    if not dry_run:
        subprocess.call(cmd, shell=True)
    else:
        click.echo(cmd)
<commit_msg>Attempt to use private IP address if public IP address is not available<commit_after># -*- coding: utf-8 -*-
import subprocess
import sys

import boto3
import botocore
import click

def format_output(instances, flag):
    """return formatted string for instance"""
    out = []
    line_format = '{0}\t{1}\t{2}\t{3}\t{4}'
    name_len = _get_max_name_len(instances) + 3
    if flag:
        line_format = '{0:<' + str(name_len) + '}{1:<10}{2:<13}{3:<16}{4:<16}'

    for i in instances:
        tag_name = get_tag_value(i.tags, 'Name')
        out.append(line_format.format(
            tag_name, i.state['Name'], i.id, i.private_ip_address, str(i.public_ip_address)))
    return out

def _get_max_name_len(instances):
    """get max length of Tag:Name"""
    # FIXME: ec2.instanceCollection doesn't have __len__
    for i in instances:
        return max([len(get_tag_value(i.tags, 'Name')) for i in instances])
    return 0

def get_tag_value(x, key):
    """Get a value from tag"""
    if x is None:
        return ''
    result = [y['Value'] for y in x if y['Key'] == key]
    if result:
        return result[0]
    return ''

@click.group()
def cli():
    """EC2 CLI group"""
    pass

@cli.command(help='List EC2 instances')
@click.argument('name', default='*')
@click.option('--list-formatted', '-l', is_flag=True)
def ls(name, list_formatted):
    """List EC2 instances"""
    ec2 = boto3.resource('ec2')
    if name == '*':
        instances = ec2.instances.filter()
    else:
        condition = {'Name': 'tag:Name', 'Values': [name]}
        instances = ec2.instances.filter(Filters=[condition])
    out = format_output(instances, list_formatted)
    click.echo('\n'.join(out))

@cli.command(help='Start EC2 instance')
@click.option('--instance-id', '-i', required=True, help='EC2 instance id')
def up(instance_id):
    """Start EC2 instance"""
    ec2 = boto3.resource('ec2')
    try:
        instance = ec2.Instance(instance_id)
        instance.start()
    except botocore.exceptions.ClientError as e:
        click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
        sys.exit(2)

@cli.command(help='Stop EC2 instance')
@click.option('--instance-id', '-i', required=True, help='EC2 instance id')
def down(instance_id):
    """Stop EC2 instance"""
    ec2 = boto3.resource('ec2')
    try:
        instance = ec2.Instance(instance_id)
        instance.stop()
    except botocore.exceptions.ClientError as e:
        click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
        sys.exit(2)

def create_ssh_command(instance_id, instance_name, username, key_file, port, gateway_instance_id):
    """Create SSH Login command string"""
    ec2 = boto3.resource('ec2')
    if instance_id is not None:
        try:
            instance = ec2.Instance(instance_id)
            hostname = instance.public_ip_address
        except botocore.exceptions.ClientError as e:
            click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
            sys.exit(2)
    elif instance_name is not None:
        try:
            condition = {'Name': 'tag:Name', 'Values': [instance_name]}
            instances = ec2.instances.filter(Filters=[condition])
            target_instances = []
            for idx, i in enumerate(instances):
                target_instances.append(i)
                tag_name = get_tag_value(i.tags, 'Name')
                click.echo('[{0}]: {1}\t{2}\t{3}\t{4}'.format(
                    idx, i.id, i.private_ip_address, i.state['Name'], tag_name))
            selected_idx = click.prompt("Please enter a valid number", type=int, default=0)
            # TODO: add validation for if selected_idx exceeds length of target_instances
            click.echo("{0} is selected.".format(selected_idx))
            instance = target_instances[selected_idx]
            if instance.public_ip_address is not None:
              hostname = instance.public_ip_address
            else:
              click.echo("Public IP address not set.  Attempting to use the private IP address.")
              hostname = instance.private_ip_address
        except botocore.exceptions.ClientError as e:
            click.echo("Invalid instance ID {0} ({1})".format(instance_id, e), err=True)
            sys.exit(2)
    # TODO: need to refactor and make it testable
    if gateway_instance_id is not None:
        gateway_instance = ec2.Instance(gateway_instance_id)
        gateway_public_ip = gateway_instance.public_ip_address
        hostname = instance.private_ip_address
        cmd = 'ssh -tt ubuntu@{0} -i {1} -p {2} ssh {3}@{4}'.format(
            gateway_public_ip, key_file, port, username, hostname)
    else:
        cmd = 'ssh {0}@{1} -i {2} -p {3}'.format(username, hostname, key_file, port)
    return cmd

@cli.command(help='SSH login to EC2 instance')
@click.option('--instance-id', '-i', default=None, help='EC2 instance id')
@click.option('--instance-name', '-n', default=None, help='EC2 instance Name Tag')
@click.option('--username', '-u', default='ubuntu', help='Login username')
@click.option('--key-file', '-k', required=True, help='SSH Key file path', type=click.Path())
@click.option('--port', '-p', help='SSH port', default=22)
@click.option('--gateway-instance-id', '-g', default=None, help='Gateway instance id')
@click.option('--dry-run', is_flag=True, default=False, help='Print SSH Login command and exist')
def ssh(instance_id, instance_name, username, key_file, port, gateway_instance_id, dry_run):
    """SSH to EC2 instance"""
    if instance_id is None and instance_name is None:
        click.echo(
            "One of --instance-id/-i or --instance-name/-n"
            " has to be specified.", err=True)
        sys.exit(2)
    elif instance_id is not None and instance_name is not None:
        click.echo(
            "Both --instance-id/-i and --instance-name/-n "
            "can't to be specified at the same time.", err=True)
        sys.exit(2)
    cmd = create_ssh_command(
        instance_id, instance_name, username, key_file, port, gateway_instance_id)
    if not dry_run:
        subprocess.call(cmd, shell=True)
    else:
        click.echo(cmd)
<|endoftext|>


<commit_before>#!/usr/bin/env python

import os
import tarfile

from contextlib import closing
from io import BytesIO

class BetterTarStream(object):
    """ Implements a tar stream from a path specified, providing at most
    `buf_size` big chunks.

    Caution: This class is specialized on the usecase and the file API only
    returns what is necessary to fullfil `requests` needs on the API.
    """

    def __init__(self, buf_size, path):
        """ Initializes a new TarStream object

        `buf_size` is the maximum buffer size to use.
        `path`     is the directory or file to create the tar stream from.
        """
        self._path = path
        self._buf_size = buf_size
        self._iterator = self._inner_next()

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        return next(self._iterator)

    def _items(self):
        """ Generator to iterate over all files and folders to be added to
        the tar archive
        """
        # Handle the case of one file only
        if os.path.isfile(self._path):
            yield self._path, os.path.basename(self._path)
        else:
            src = os.path.normpath(self._path)
            src_prefix = os.path.dirname(src)
            for root, dirs, files in os.walk(src, topdown=True):
                rel_root = os.path.relpath(root, src_prefix)
                for name in dirs + files:
                    item = os.path.join(rel_root, name)
                    yield os.path.join(src_prefix, item), item

    def _inner_next(self):
        """ Generator that iterates over all streams and yields buffer chunks
        in the size of the buf_size parameter
        """
        for stream in self._streams():
            while True:
                buf = stream.read(self._buf_size)
                if not buf:
                    stream.close()
                    break
                yield buf
        yield bytes()

    def _streams(self):
        """ Generator that creates streams representing the tar format """
        tfile = tarfile.open(fileobj=BytesIO(), bufsize=self._buf_size,
                             mode='w|')
        with closing(tfile):
            for path, arcname in self._items():
                info = tfile.gettarinfo(path, arcname)
                # Yield the header as a stream
                yield BytesIO(info.tobuf())

                if info.isreg():
                    # Yield the file stream
                    yield open(path, 'rb')

                padding = info.size % tarfile.BLOCKSIZE
                if padding:
                    padding = tarfile.BLOCKSIZE - padding
                if padding:
                    # Yield the padding as a stream if necessary
                    yield BytesIO(padding * tarfile.NUL)

        # Yield the tar finalization marker
        yield BytesIO((tarfile.BLOCKSIZE * 2) * tarfile.NUL)
<commit_msg>A little bit cleaner implementation for the tar stream buffer generator<commit_after>#!/usr/bin/env python

import os
import tarfile

from contextlib import closing
from io import BytesIO

class BetterTarStream(object):
    """ Implements a tar stream from a path specified, providing at most
    `buf_size` big chunks.

    Caution: This class is specialized on the usecase and the file API only
    returns what is necessary to fullfil `requests` needs on the API.
    """

    def __init__(self, buf_size, path):
        """ Initializes a new TarStream object

        `buf_size` is the maximum buffer size to use.
        `path`     is the directory or file to create the tar stream from.
        """
        self._path = path
        self._buf_size = buf_size
        self._iterator = self._inner_next()

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        return next(self._iterator)

    def _items(self):
        """ Generator to iterate over all files and folders to be added to
        the tar archive
        """
        # Handle the case of one file only
        if os.path.isfile(self._path):
            yield self._path, os.path.basename(self._path)
        else:
            src = os.path.normpath(self._path)
            src_prefix = os.path.dirname(src)
            for root, dirs, files in os.walk(src, topdown=True):
                rel_root = os.path.relpath(root, src_prefix)
                for name in dirs + files:
                    item = os.path.join(rel_root, name)
                    yield os.path.join(src_prefix, item), item

    def _inner_next(self):
        """ Generator that iterates over all streams and yields buffer chunks
        in the size of the buf_size parameter
        """
        for stream in self._streams():
            with closing(stream):
                while True:
                    buf = stream.read(self._buf_size)
                    if not buf:
                        break
                    yield buf

    def _streams(self):
        """ Generator that creates streams representing the tar format """
        tfile = tarfile.open(fileobj=BytesIO(), bufsize=self._buf_size,
                             mode='w|')
        with closing(tfile):
            for path, arcname in self._items():
                info = tfile.gettarinfo(path, arcname)
                # Yield the header as a stream
                yield BytesIO(info.tobuf())

                if info.isreg():
                    # Yield the file stream
                    yield open(path, 'rb')

                padding = info.size % tarfile.BLOCKSIZE
                if padding:
                    padding = tarfile.BLOCKSIZE - padding
                if padding:
                    # Yield the padding as a stream if necessary
                    yield BytesIO(padding * tarfile.NUL)

        # Yield the tar finalization marker
        yield BytesIO((tarfile.BLOCKSIZE * 2) * tarfile.NUL)
<|endoftext|>


<commit_before>"""
A Spawner for JupyterHub that runs each user's server in a separate docker container
"""
import itertools
import os
from textwrap import dedent
from concurrent.futures import ThreadPoolExecutor

import docker
from tornado import gen

from jupyterhub.spawner import Spawner
from IPython.utils.traitlets import (
    Dict,
    Unicode,
)

class DockerSpawner(Spawner):

    _executor = None
    @property
    def executor(self):
        """single global executor"""
        cls = self.__class__
        if cls._executor is None:
            cls._executor = ThreadPoolExecutor(1)
        return cls._executor

    _client = None
    @property
    def client(self):
        """single global client instance"""
        cls = self.__class__
        if cls._client is None:
            docker_host = os.environ.get('DOCKER_HOST', 'unix://var/run/docker.sock')
            cls._client = docker.Client(base_url=docker_host)
        return cls._client

    container_id = Unicode()
    container_ip = Unicode('127.0.0.1', config=True)
    container_image = Unicode("jupyter/singleuser", config=True)
    container_prefix = Unicode(
        "jupyter",
        config=True,
        help=dedent(
            """
            Prefix for container names. The full container name for a particular
            user will be <prefix>-<username>.
            """
        )
    )

    volumes = Dict(
        config=True,
        help=dedent(
            """
            Map from host file/directory to container file/directory.
            Volumes specified here will be read/write in the container.
            """
        )
    )
    read_only_volumes = Dict(
        config=True,
        help=dedent(
            """
            Map from host file/directory to container file/directory.
            Volumes specified here will be read-only in the container.
            """
        )
    )

    @property
    def volume_mount_points(self):
        """
        Volumes are declared in docker-py in two stages.  First, you declare
        all the locations where you're going to mount volumes when you call
        create_container.

        Returns a list of all the values in self.volumes or
        self.read_only_volumes.
        """
        return list(
            itertools.chain(
                self.volumes.values(),
                self.read_only_volumes.values(),
            )
        )

    @property
    def volume_binds(self):
        """
        The second half of declaring a volume with docker-py happens when you
        actually call start().  The required format is a dict of dicts that
        looks like:

        {
            host_location: {'bind': container_location, 'ro': True}
        }
        """
        volumes = {
            key: {'bind': value, 'ro': False}
            for key, value in self.volumes.items()
        }
        ro_volumes = {
            key: {'bind': value, 'ro': True}
            for key, value in self.read_only_volumes.items()
        }
        volumes.update(ro_volumes)
        return volumes

    @property
    def container_name(self):
        return "{}-{}".format(self.container_prefix, self.user.name)

    def load_state(self, state):
        super(DockerSpawner, self).load_state(state)
        self.container_id = state.get('container_id', '')

    def get_state(self):
        state = super(DockerSpawner, self).get_state()
        if self.container_id:
            state['container_id'] = self.container_id
        return state

    def _env_keep_default(self):
        """Don't inherit any env from the parent process"""
        return []

    def _env_default(self):
        env = super(DockerSpawner, self)._env_default()
        env.update(dict(
            JPY_USER=self.user.name,
            JPY_COOKIE_NAME=self.user.server.cookie_name,
            JPY_BASE_URL=self.user.server.base_url,
            JPY_HUB_PREFIX=self.hub.server.base_url,
            JPY_HUB_API_URL=self.hub.api_url,
        ))
        return env

    def _docker(self, method, *args, **kwargs):
        """wrapper for calling docker methods

        to be passed to ThreadPoolExecutor
        """
        m = getattr(self.client, method)
        return m(*args, **kwargs)

    def docker(self, method, *args, **kwargs):
        """Call a docker method in a background thread

        returns a Future
        """
        return self.executor.submit(self._docker, method, *args, **kwargs)

    @gen.coroutine
    def poll(self):
        """Check for my id in `docker ps`"""
        container = yield self.get_container()
        if not container:
            self.log.warn("container not found")
            raise gen.Return(0)
        status = container['Status']
        self.log.debug("Container %s status: %s", self.container_id[:7], status)
        parts = status.split()
        # status examples:
        # 'Exited (127) 12 days ago'
        # 'Up 4 days'
        # 'Up 4 days (Paused)' <- TODO: this isn't handled
        # e.g.
        if parts[0] == 'Up':
            raise gen.Return(None)
        elif parts[0] == 'Exited':
            raise gen.Return(int(parts[1][1:-1]))
        else:
            raise ValueError("Unhandled status: '%s'" % status)

    @gen.coroutine
    def get_container(self):
        self.log.debug("Getting container '%s'", self.container_name)
        containers = yield self.docker('containers', all=True)
        for c in containers:
            if "/{}".format(self.container_name) in c['Names']:
                self.container_id = c['Id']
                raise gen.Return(c)
        self.log.info("Container '%s' is gone", self.container_name)
        # my container is gone, forget my id
        self.container_id = ''

    @gen.coroutine
    def start(self, image=None, **extra_create_kwargs):
        """Start the single-user server in a docker container. You can override
        the default parameters passed to `create_container` through the
        `extra_create_kwargs` dictionary.

        """
        container = yield self.get_container()
        if container is None:
            image = image or self.container_image

            # build the dictionary of keyword arguments for create_container
            create_kwargs = dict(
                image=image,
                environment=self.env,
                volumes=self.volume_mount_points,
                name=self.container_name)
            create_kwargs.update(extra_create_kwargs)

            # create the container
            resp = yield self.docker('create_container', **create_kwargs)
            self.container_id = resp['Id']
            self.log.info(
                "Created container '%s' (id: %s) from image %s",
                self.container_name, self.container_id[:7], image)

        else:
            self.log.info(
                "Found existing container '%s' (id: %s)",
                self.container_name, self.container_id[:7])

        # TODO: handle unpause
        self.log.info(
            "Starting container '%s' (id: %s)",
            self.container_name, self.container_id[:7])
        yield self.docker(
            'start',
            self.container_id,
            binds=self.volume_binds,
            port_bindings={8888: (self.container_ip,)}
        )
        resp = yield self.docker('port', self.container_id, 8888)
        self.user.server.port = resp[0]['HostPort']

    @gen.coroutine
    def stop(self, now=False):
        """Stop the container

        Consider using pause/unpause when docker-py adds support
        """
        self.log.info(
            "Stopping container %s' (id: %s)",
            self.container_name, self.container_id[:7])
        yield self.docker('stop', self.container_id)
        self.clear_state()
<commit_msg>MAINT: Remove stray quote from log message.<commit_after>"""
A Spawner for JupyterHub that runs each user's server in a separate docker container
"""
import itertools
import os
from textwrap import dedent
from concurrent.futures import ThreadPoolExecutor

import docker
from tornado import gen

from jupyterhub.spawner import Spawner
from IPython.utils.traitlets import (
    Dict,
    Unicode,
)

class DockerSpawner(Spawner):

    _executor = None
    @property
    def executor(self):
        """single global executor"""
        cls = self.__class__
        if cls._executor is None:
            cls._executor = ThreadPoolExecutor(1)
        return cls._executor

    _client = None
    @property
    def client(self):
        """single global client instance"""
        cls = self.__class__
        if cls._client is None:
            docker_host = os.environ.get('DOCKER_HOST', 'unix://var/run/docker.sock')
            cls._client = docker.Client(base_url=docker_host)
        return cls._client

    container_id = Unicode()
    container_ip = Unicode('127.0.0.1', config=True)
    container_image = Unicode("jupyter/singleuser", config=True)
    container_prefix = Unicode(
        "jupyter",
        config=True,
        help=dedent(
            """
            Prefix for container names. The full container name for a particular
            user will be <prefix>-<username>.
            """
        )
    )

    volumes = Dict(
        config=True,
        help=dedent(
            """
            Map from host file/directory to container file/directory.
            Volumes specified here will be read/write in the container.
            """
        )
    )
    read_only_volumes = Dict(
        config=True,
        help=dedent(
            """
            Map from host file/directory to container file/directory.
            Volumes specified here will be read-only in the container.
            """
        )
    )

    @property
    def volume_mount_points(self):
        """
        Volumes are declared in docker-py in two stages.  First, you declare
        all the locations where you're going to mount volumes when you call
        create_container.

        Returns a list of all the values in self.volumes or
        self.read_only_volumes.
        """
        return list(
            itertools.chain(
                self.volumes.values(),
                self.read_only_volumes.values(),
            )
        )

    @property
    def volume_binds(self):
        """
        The second half of declaring a volume with docker-py happens when you
        actually call start().  The required format is a dict of dicts that
        looks like:

        {
            host_location: {'bind': container_location, 'ro': True}
        }
        """
        volumes = {
            key: {'bind': value, 'ro': False}
            for key, value in self.volumes.items()
        }
        ro_volumes = {
            key: {'bind': value, 'ro': True}
            for key, value in self.read_only_volumes.items()
        }
        volumes.update(ro_volumes)
        return volumes

    @property
    def container_name(self):
        return "{}-{}".format(self.container_prefix, self.user.name)

    def load_state(self, state):
        super(DockerSpawner, self).load_state(state)
        self.container_id = state.get('container_id', '')

    def get_state(self):
        state = super(DockerSpawner, self).get_state()
        if self.container_id:
            state['container_id'] = self.container_id
        return state

    def _env_keep_default(self):
        """Don't inherit any env from the parent process"""
        return []

    def _env_default(self):
        env = super(DockerSpawner, self)._env_default()
        env.update(dict(
            JPY_USER=self.user.name,
            JPY_COOKIE_NAME=self.user.server.cookie_name,
            JPY_BASE_URL=self.user.server.base_url,
            JPY_HUB_PREFIX=self.hub.server.base_url,
            JPY_HUB_API_URL=self.hub.api_url,
        ))
        return env

    def _docker(self, method, *args, **kwargs):
        """wrapper for calling docker methods

        to be passed to ThreadPoolExecutor
        """
        m = getattr(self.client, method)
        return m(*args, **kwargs)

    def docker(self, method, *args, **kwargs):
        """Call a docker method in a background thread

        returns a Future
        """
        return self.executor.submit(self._docker, method, *args, **kwargs)

    @gen.coroutine
    def poll(self):
        """Check for my id in `docker ps`"""
        container = yield self.get_container()
        if not container:
            self.log.warn("container not found")
            raise gen.Return(0)
        status = container['Status']
        self.log.debug("Container %s status: %s", self.container_id[:7], status)
        parts = status.split()
        # status examples:
        # 'Exited (127) 12 days ago'
        # 'Up 4 days'
        # 'Up 4 days (Paused)' <- TODO: this isn't handled
        # e.g.
        if parts[0] == 'Up':
            raise gen.Return(None)
        elif parts[0] == 'Exited':
            raise gen.Return(int(parts[1][1:-1]))
        else:
            raise ValueError("Unhandled status: '%s'" % status)

    @gen.coroutine
    def get_container(self):
        self.log.debug("Getting container '%s'", self.container_name)
        containers = yield self.docker('containers', all=True)
        for c in containers:
            if "/{}".format(self.container_name) in c['Names']:
                self.container_id = c['Id']
                raise gen.Return(c)
        self.log.info("Container '%s' is gone", self.container_name)
        # my container is gone, forget my id
        self.container_id = ''

    @gen.coroutine
    def start(self, image=None, **extra_create_kwargs):
        """Start the single-user server in a docker container. You can override
        the default parameters passed to `create_container` through the
        `extra_create_kwargs` dictionary.

        """
        container = yield self.get_container()
        if container is None:
            image = image or self.container_image

            # build the dictionary of keyword arguments for create_container
            create_kwargs = dict(
                image=image,
                environment=self.env,
                volumes=self.volume_mount_points,
                name=self.container_name)
            create_kwargs.update(extra_create_kwargs)

            # create the container
            resp = yield self.docker('create_container', **create_kwargs)
            self.container_id = resp['Id']
            self.log.info(
                "Created container '%s' (id: %s) from image %s",
                self.container_name, self.container_id[:7], image)

        else:
            self.log.info(
                "Found existing container '%s' (id: %s)",
                self.container_name, self.container_id[:7])

        # TODO: handle unpause
        self.log.info(
            "Starting container '%s' (id: %s)",
            self.container_name, self.container_id[:7])
        yield self.docker(
            'start',
            self.container_id,
            binds=self.volume_binds,
            port_bindings={8888: (self.container_ip,)}
        )
        resp = yield self.docker('port', self.container_id, 8888)
        self.user.server.port = resp[0]['HostPort']

    @gen.coroutine
    def stop(self, now=False):
        """Stop the container

        Consider using pause/unpause when docker-py adds support
        """
        self.log.info(
            "Stopping container %s (id: %s)",
            self.container_name, self.container_id[:7])
        yield self.docker('stop', self.container_id)
        self.clear_state()
<|endoftext|>


<commit_before>#!/usr/bin/python

# http://repo1.maven.org/maven2/org/mongodb/mongo-java-driver/

from xml.dom.minidom import parseString
import datetime
import os
import sys
import shutil
import subprocess
import time

if len( sys.argv ) == 1:
    print "Usage: mavenPush.py VERSION [PUBDIR]"
    print "VERSION - version you want to publish"
    print "PUBDIR - directory to publish to. default = /ebs/maven/"
    sys.exit()

version = sys.argv[1]

if len( sys.argv ) > 2:
    root = os.path.expanduser( sys.argv[2] )
else:
    root = "/ebs/maven/"

p = subprocess.Popen( [ "/usr/bin/ant" , "clean" ] , stdout=subprocess.PIPE ).communicate()
p = subprocess.Popen( [ "/usr/bin/ant" , "alljars" ] , stdout=subprocess.PIPE ).communicate()

if p[0].find( "SUCCESSFUL" ) < 0:
    print( p[0] )
    print( p[1] )
    raise( "build failed" )

def build_metadata_xml(path, artifactid):
    groupid = os.path.split(path)[1]
    xml = '<metadata>'
    xml += '<groupId>org.%s</groupId>' % (groupid,)
    xml += '<artifactId>%s</artifactId>' % (artifactid,)
    xml += '<versioning><versions>'

    entries = os.listdir(path)
    for entry in entries:
        if os.path.isdir(os.path.join(path, entry)):
            xml += '<version>%s</version>' % (entry,)
    xml += '</versions>'
    xml += '<lastUpdated>%s</lastUpdated>' % (int(time.time()*1000),)
    xml += '</versioning></metadata>'

    doc = parseString(xml)
    return doc.documentElement.toprettyxml()

def go( pkgName, shortName , longName ):
    dir = root + pkgName + longName + "/" + version
    if not os.path.exists( dir ):
        os.makedirs( dir )

    fileRoot = dir + "/" + longName + "-" + version

    for x in [ "" , "-sources" , "-javadoc" ]:
        shutil.copy2( shortName + x + ".jar" , fileRoot + x + ".jar" )

    pom = open( "maven-" + shortName + ".xml" , "r" ).read()
    pom = pom.replace( "$VERSION" , version )

    out = open( fileRoot + ".pom" , 'w' )
    out.write( pom )
    out.close()

    p = subprocess.Popen( [ "sha1sum" , fileRoot + ".jar" ] , stdout=subprocess.PIPE ).communicate()
    sha1 = p[0].split( ' ' )[0]
    out = open( fileRoot + ".jar.sha1" , 'w' )
    out.write( sha1 )
    out.close()

    out = open( os.path.join(os.path.split(dir)[0], "maven-metadata.xml") , 'w' )
    out.write( build_metadata_xml(os.path.split(dir)[0], longName) )
    out.close()

go( "/org/mongodb/" , "mongo" , "mongo-java-driver" )
go( "/org/bson/" , "bson" , "bson" )
<commit_msg>moved to ant tasks.<commit_after><|endoftext|>


<commit_before># -*- coding: UTF-8 -*-

__version__ = '0.0.1'

class KEnum(int):
    pass

def make(name, names, start=0):
    """
    Create an enum names ``name``, with a given list of names for each member.
    """
    klass = type(name, (KEnum,), {})

    for k, v in zip(names, range(start, len(names) + start)):
        setattr(klass, k, klass(v))

    return klass
<commit_msg>PEP8 compliance<commit_after># -*- coding: UTF-8 -*-

__version__ = '0.0.1'

class KEnum(int):
    pass

def make(name, names, start=0):
    """
    Create an enum names ``name``, with a given list of names for each member.
    """
    klass = type(name, (KEnum,), {})

    for k, v in zip(names, range(start, len(names) + start)):
        setattr(klass, k, klass(v))

    return klass
<|endoftext|>


<commit_before><commit_msg><commit_after><|endoftext|>


<commit_before>import unittest
import bdkd.datastore
import glob, os, shutil

FIXTURES = os.path.join(os.path.dirname(__file__), '..', '..', 'fixtures')

class RepositoryTest(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        super(self.__class__, self).__init__(*args, **kwargs)
        self.repository = bdkd.datastore.repository("bdkd-datastore-test")
        self.assertTrue(self.repository)
        self.assertTrue(self.repository.get_bucket())
        self.resources = dict(
                single=self._create_single_file_fixture(),
                multi=self._create_multi_file_fixture(),
                shapefile=self._create_shapefile_fixture(),
                )

    def setUp(self):
        # Start with the test bucket empty
        for key in self.repository.get_bucket().list():
            key.delete()
        self._clear_local()

    def test_single_file_resource(self):
        resource = self.resources.get('single')
        self.assertTrue(resource)
        # Save the resource to the repository
        self.repository.save(resource)
        # Check that the repository was set properly
        self.assertTrue(resource.repository == self.repository)
        # Check that the Resource was copied to the local cache
        resource_path = os.path.join(self.repository.local_cache,
                'resources', resource.name)
        self.assertTrue(os.path.exists(resource_path))
        # Check that S3 repository has two keys in it (resource + one file)
        self._check_bucket_count('', 2)
        # Check that the cache has two files in it
        self._check_file_count(self.repository.local_cache, 2)

    def test_name_collisions(self):
        resource = self.resources.get('single')
        other_resource = self.resources.get('multi')

        resource.name = 'a/b'
        # Save the resource to the repository
        self.repository.save(resource)

        # Conflicts with the shorter name 'a/b'
        with self.assertRaises(ValueError):
            other_resource.name = 'a/b/c'
            self.repository.save(other_resource)

        # Conflicts with the longer name 'a/b'
        with self.assertRaises(ValueError):
            other_resource.name = 'a'
            self.repository.save(other_resource)

    def test_multi_file_resource(self):
        resource = self.resources.get('multi')
        self.assertTrue(resource)
        # Save the resource to the repository
        self.repository.save(resource)
        # Check that the repository was set properly
        self.assertTrue(resource.repository == self.repository)
        # Check that the Resource was copied to the local cache
        resource_path = os.path.join(self.repository.local_cache,
                'resources', resource.name)
        self.assertTrue(os.path.exists(resource_path))
        # Check that S3 repository has two keys in it (resource + one file)
        self._check_bucket_count('', 3)
        # Check that the cache has two files in it
        self._check_file_count(self.repository.local_cache, 3)

    def test_shapefile_resource(self):
        resource = self.resources.get('shapefile')
        self.assertTrue(resource)
        self.repository.save(resource)
        # The shapefile consists of five files (plus one resource)
        self._check_bucket_count('', 6)
        self._check_file_count(self.repository.local_cache, 6)

    def test_unified_resource(self):
        """ Unified: all resource parts fetched together. """
        resource = self.resources.get('shapefile')
        self.assertTrue(resource)
        self.repository.save(resource)
        self._clear_local()
        shp = resource.file_ending('.shp')
        shp.local_path()
        # The shapefile consists of five files (plus one resource)
        self._check_file_count(self.repository.local_cache, 6)

    def test_remote_resource(self):
        resource = bdkd.datastore.Resource.new('Caltech/Continuously Closing Plate Polygons',
                'http://www.gps.caltech.edu/~gurnis/GPlates/Caltech_Global_20101129.tar.gz')
        self.assertTrue(resource.files[0].meta('etag'))
        self.repository.save(resource)
        paths = resource.local_paths()
        self.assertEquals(len(paths), 1)
        self.assertEquals(bdkd.datastore.checksum(paths[0]), '412f9bfd1fc6c6ba1b5fc8bc450fef61')

    def test_refresh_cache(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # Clear the local cache, refresh the resource and see if the cache is restored
        self._clear_local()
        self._check_file_count(self.repository.local_cache, 0)
        self.repository.refresh_resource(resource)
        self._check_file_count(self.repository.local_cache, 2)

    def test_list_resources(self):
        # Starts with nothing there
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 0)
        # One resource
        resource = self.resources.get('single')
        self.repository.save(resource)
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 1)
        self.assertEquals(resource_names[0], resource.name)
        # More than one
        shapefile = self.resources.get('shapefile')
        self.repository.save(shapefile)
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 2)
        # Listing by pseudopath
        resource_names = self.repository.list(resource.name)
        self.assertEquals(len(resource_names), 1)
        # Listing by nonexistent pseudopath
        resource_names = self.repository.list('foo/bar')
        self.assertEquals(len(resource_names), 0)

    def test_edit_resource(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # After saving, the Resource should not be in editing mode.
        self.assertFalse(resource.is_edit)
        # The Resource path should be a file in the cache directory that is not writable.
        self.assertTrue(resource.path.startswith(self.repository.local_cache))
        self.assertFalse(os.access(resource.path, os.W_OK))
        self._check_file_count(self.repository.working, 0)
        # Now set resource to editing mode
        self.repository.edit_resource(resource)
        self.assertTrue(resource.is_edit)
        # There should be a writable file in the repository's working path
        self.assertTrue(resource.path.startswith(self.repository.working))
        self.assertTrue(os.access(resource.path, os.W_OK))
        self._check_file_count(self.repository.working, 2)

    def test_delete_resource(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # Before delete: two files
        self._check_bucket_count('', 2)
        self._check_file_count(self.repository.local_cache, 2)

        self.repository.delete(resource)
        # After delete: no files
        self._check_bucket_count('', 0)
        self._check_file_count(self.repository.local_cache, 0)

    def _clear_local(self):
        for tmp_path in [ self.repository.local_cache, self.repository.working ]:
            if tmp_path and tmp_path.startswith('/var/tmp'):
                if os.path.exists(tmp_path):
                    shutil.rmtree(tmp_path)

    def _check_file_count(self, dir_path, expected):
        file_count = 0
        for dirpath, dirnames, filenames in os.walk(dir_path):
            file_count += len(filenames)
        self.assertTrue(file_count == expected)

    def _check_bucket_count(self, pseudopath, expected):
        # Check that S3 repository has two keys in it
        key_count = 0
        for key in self.repository.get_bucket().list(pseudopath):
            key_count += 1
        self.assertTrue(key_count == expected)

    def _create_single_file_fixture(self):
        path = os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                'Seton_etal_ESR2012_Coastlines_2012.1.gpmlz')
        resource = bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Seton',
                path)
        return resource

    def _create_multi_file_fixture(self):
        return bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Seton',
                [
                    os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                        'Seton_etal_ESR2012_Coastlines_2012.1.gpmlz'),
                    os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                        'Shapefile', 'Seton_etal_ESR2012_Coastlines_2012.1.shp')
                    ])

    def _create_shapefile_fixture(self):
        shapefile_dir = os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines', 'Shapefile')
        shapefile_parts = glob.glob(os.path.join(shapefile_dir, '*.*'))
        files = []
        for part in shapefile_parts:
            file_meta = dict(path=os.path.expanduser(part))
            if part.endswith('shp'):
                file_meta['type'] = 'ESRI shapefile'
                files.insert(0, file_meta)
            else:
                files.append(file_meta)
        resource = bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Shapefile/Seton',
                files, unified=True)
        return resource

if __name__ == '__main__':
    unittest.main()

<commit_msg>BDKD-45: Fixed Python 2.6 incompatibility with the use of assertRaises().<commit_after>import unittest
import bdkd.datastore
import glob, os, shutil

FIXTURES = os.path.join(os.path.dirname(__file__), '..', '..', 'fixtures')

class RepositoryTest(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        super(self.__class__, self).__init__(*args, **kwargs)
        self.repository = bdkd.datastore.repository("bdkd-datastore-test")
        self.assertTrue(self.repository)
        self.assertTrue(self.repository.get_bucket())
        self.resources = dict(
                single=self._create_single_file_fixture(),
                multi=self._create_multi_file_fixture(),
                shapefile=self._create_shapefile_fixture(),
                )

    def setUp(self):
        # Start with the test bucket empty
        for key in self.repository.get_bucket().list():
            key.delete()
        self._clear_local()

    def test_single_file_resource(self):
        resource = self.resources.get('single')
        self.assertTrue(resource)
        # Save the resource to the repository
        self.repository.save(resource)
        # Check that the repository was set properly
        self.assertTrue(resource.repository == self.repository)
        # Check that the Resource was copied to the local cache
        resource_path = os.path.join(self.repository.local_cache,
                'resources', resource.name)
        self.assertTrue(os.path.exists(resource_path))
        # Check that S3 repository has two keys in it (resource + one file)
        self._check_bucket_count('', 2)
        # Check that the cache has two files in it
        self._check_file_count(self.repository.local_cache, 2)

    def test_name_collisions(self):
        resource = self.resources.get('single')
        other_resource = self.resources.get('multi')

        resource.name = 'a/b'
        # Save the resource to the repository
        self.repository.save(resource)

        # Conflicts with the shorter name 'a/b'
        other_resource.name = 'a/b/c'
        self.assertRaises(ValueError, self.repository.save, other_resource)

        # Conflicts with the longer name 'a/b'
        other_resource.name = 'a'
        self.assertRaises(ValueError, self.repository.save, other_resource)

    def test_multi_file_resource(self):
        resource = self.resources.get('multi')
        self.assertTrue(resource)
        # Save the resource to the repository
        self.repository.save(resource)
        # Check that the repository was set properly
        self.assertTrue(resource.repository == self.repository)
        # Check that the Resource was copied to the local cache
        resource_path = os.path.join(self.repository.local_cache,
                'resources', resource.name)
        self.assertTrue(os.path.exists(resource_path))
        # Check that S3 repository has two keys in it (resource + one file)
        self._check_bucket_count('', 3)
        # Check that the cache has two files in it
        self._check_file_count(self.repository.local_cache, 3)

    def test_shapefile_resource(self):
        resource = self.resources.get('shapefile')
        self.assertTrue(resource)
        self.repository.save(resource)
        # The shapefile consists of five files (plus one resource)
        self._check_bucket_count('', 6)
        self._check_file_count(self.repository.local_cache, 6)

    def test_unified_resource(self):
        """ Unified: all resource parts fetched together. """
        resource = self.resources.get('shapefile')
        self.assertTrue(resource)
        self.repository.save(resource)
        self._clear_local()
        shp = resource.file_ending('.shp')
        shp.local_path()
        # The shapefile consists of five files (plus one resource)
        self._check_file_count(self.repository.local_cache, 6)

    def test_remote_resource(self):
        resource = bdkd.datastore.Resource.new('Caltech/Continuously Closing Plate Polygons',
                'http://www.gps.caltech.edu/~gurnis/GPlates/Caltech_Global_20101129.tar.gz')
        self.assertTrue(resource.files[0].meta('etag'))
        self.repository.save(resource)
        paths = resource.local_paths()
        self.assertEquals(len(paths), 1)
        self.assertEquals(bdkd.datastore.checksum(paths[0]), '412f9bfd1fc6c6ba1b5fc8bc450fef61')

    def test_refresh_cache(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # Clear the local cache, refresh the resource and see if the cache is restored
        self._clear_local()
        self._check_file_count(self.repository.local_cache, 0)
        self.repository.refresh_resource(resource)
        self._check_file_count(self.repository.local_cache, 2)

    def test_list_resources(self):
        # Starts with nothing there
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 0)
        # One resource
        resource = self.resources.get('single')
        self.repository.save(resource)
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 1)
        self.assertEquals(resource_names[0], resource.name)
        # More than one
        shapefile = self.resources.get('shapefile')
        self.repository.save(shapefile)
        resource_names = self.repository.list()
        self.assertEquals(len(resource_names), 2)
        # Listing by pseudopath
        resource_names = self.repository.list(resource.name)
        self.assertEquals(len(resource_names), 1)
        # Listing by nonexistent pseudopath
        resource_names = self.repository.list('foo/bar')
        self.assertEquals(len(resource_names), 0)

    def test_edit_resource(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # After saving, the Resource should not be in editing mode.
        self.assertFalse(resource.is_edit)
        # The Resource path should be a file in the cache directory that is not writable.
        self.assertTrue(resource.path.startswith(self.repository.local_cache))
        self.assertFalse(os.access(resource.path, os.W_OK))
        self._check_file_count(self.repository.working, 0)
        # Now set resource to editing mode
        self.repository.edit_resource(resource)
        self.assertTrue(resource.is_edit)
        # There should be a writable file in the repository's working path
        self.assertTrue(resource.path.startswith(self.repository.working))
        self.assertTrue(os.access(resource.path, os.W_OK))
        self._check_file_count(self.repository.working, 2)

    def test_delete_resource(self):
        resource = self.resources.get('single')
        self.repository.save(resource)
        # Before delete: two files
        self._check_bucket_count('', 2)
        self._check_file_count(self.repository.local_cache, 2)

        self.repository.delete(resource)
        # After delete: no files
        self._check_bucket_count('', 0)
        self._check_file_count(self.repository.local_cache, 0)

    def _clear_local(self):
        for tmp_path in [ self.repository.local_cache, self.repository.working ]:
            if tmp_path and tmp_path.startswith('/var/tmp'):
                if os.path.exists(tmp_path):
                    shutil.rmtree(tmp_path)

    def _check_file_count(self, dir_path, expected):
        file_count = 0
        for dirpath, dirnames, filenames in os.walk(dir_path):
            file_count += len(filenames)
        self.assertTrue(file_count == expected)

    def _check_bucket_count(self, pseudopath, expected):
        # Check that S3 repository has two keys in it
        key_count = 0
        for key in self.repository.get_bucket().list(pseudopath):
            key_count += 1
        self.assertTrue(key_count == expected)

    def _create_single_file_fixture(self):
        path = os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                'Seton_etal_ESR2012_Coastlines_2012.1.gpmlz')
        resource = bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Seton',
                path)
        return resource

    def _create_multi_file_fixture(self):
        return bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Seton',
                [
                    os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                        'Seton_etal_ESR2012_Coastlines_2012.1.gpmlz'),
                    os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines',
                        'Shapefile', 'Seton_etal_ESR2012_Coastlines_2012.1.shp')
                    ])

    def _create_shapefile_fixture(self):
        shapefile_dir = os.path.join(FIXTURES, 'FeatureCollections', 'Coastlines', 'Shapefile')
        shapefile_parts = glob.glob(os.path.join(shapefile_dir, '*.*'))
        files = []
        for part in shapefile_parts:
            file_meta = dict(path=os.path.expanduser(part))
            if part.endswith('shp'):
                file_meta['type'] = 'ESRI shapefile'
                files.insert(0, file_meta)
            else:
                files.append(file_meta)
        resource = bdkd.datastore.Resource.new('FeatureCollections/Coastlines/Shapefile/Seton',
                files, unified=True)
        return resource

if __name__ == '__main__':
    unittest.main()

<|endoftext|>