bootstrap/bootstrap/__init__.py

"""Interactively set SSH keys on remote Git servers."""

from __future__ import print_function
from getpass import getpass
from os import environ
from os.path import join
from platform import node, system

from requests import ConnectionError, get, post
from requests.auth import HTTPBasicAuth
from requests.compat import urlparse

try:
    input = raw_input
except NameError:
    pass


class BootstrapError(Exception):
    """Bootstrap Exception."""
    pass


def bootstrap_error(response):
    """Create a BootstrapError from a Response."""
    return BootstrapError('%s %s' % (response.status_code, response.reason))


def agree(question, default='Y'):
    """Prompt user to answer a yes/no question."""
    valid = {
        'y': True,
        'Y': True,
        'yes': True,
        'n': False,
        'N': False,
        'no': False,
        '': {'Y': True,
             'N': False}[default]
    }
    answer = input('%s [%s]? ' % (question, {'Y': 'Y/n', 'N': 'y/N'}[default]))
    try:
        return valid[answer]
    except KeyError:
        print('invalid input: %s' % answer)
        return agree(question)


def get_url(service, default):
    """Get URL."""
    url = input('%s URL%s: ' % (service, ' (%s)' % default if default else ''))
    if url == '':
        if default:
            return default
        else:
            print('invalid input: %s' % url)
            return get_url(service, default)
    return url


def get_username_password(service):
    """Get username/password."""
    username = input('%s username: ' % service)
    password = getpass('%s password: ' % service)
    return (username, password)


def get_local_key():
    """Get local SSH key."""
    if system() == 'Windows':
        home = environ['userprofile']
    else:
        home = environ['HOME']
    with open(join(home, '.ssh', 'id_rsa.pub'), 'r') as key_file:
        return key_file.read().rstrip()


def key_exists(keys, local_key):
    """Check if local SSH key is already set."""
    for key in keys:
        if local_key.startswith(key['key']):
            return True
    return False


def set_github_ssh_key():
    """Set GitHub SSH key."""
    url = urlparse(get_url('GitHub', 'https://github.com'))
    keys_url = '%s://api.%s/user/keys' % (url.scheme, url.netloc)
    username, password = get_username_password('GitHub')
    auth = HTTPBasicAuth(username, password)
    response = get(keys_url, auth=auth)
    if response.status_code != 200:
        raise bootstrap_error(response)
    keys = response.json()
    local_key = get_local_key()
    if not key_exists(keys, local_key):
        response = post(
            keys_url, auth=auth, json={'title': node(),
                                       'key': local_key})
        if response.status_code != 201:
            raise bootstrap_error(response)


def set_gitlab_ssh_key():
    """Set GitLab SSH key."""
    api_url = '%s/api/v4' % get_url('GitLab', 'https://gitlab.com')
    session_url = '%s/session' % api_url
    keys_url = '%s/user/keys' % api_url
    username, password = get_username_password('GitLab')
    response = post(session_url, {'login': username, 'password': password})
    if response.status_code != 201:
        raise bootstrap_error(response)
    auth = {'Private-Token': response.json()['private_token']}
    response = get(keys_url, headers=auth)
    if response.status_code != 200:
        raise bootstrap_error(response)
    keys = response.json()
    local_key = get_local_key()
    if not key_exists(keys, local_key):
        response = post(
            keys_url, headers=auth, json={'title': node(),
                                          'key': local_key})
        if response.status_code != 201:
            raise bootstrap_error(response)


def set_bitbucket_cloud_ssh_key():
    """Set BitBucket Cloud SSH key."""
    username, password = get_username_password('BitBucket Cloud')
    keys_url = 'https://api.bitbucket.org/1.0/users/%s/ssh-keys' % username
    auth = HTTPBasicAuth(username, password)
    response = get(keys_url, auth=auth)
    if response.status_code != 200:
        raise bootstrap_error(response)
    keys = response.json()
    local_key = get_local_key()
    if not key_exists(keys, local_key):
        response = post(
            keys_url, auth=auth, data={'label': node(),
                                       'key': local_key})
        if response.status_code != 200:
            raise bootstrap_error(response)


def set_gogs_ssh_key():
    """Set Gogs SSH key."""
    keys_url = '%s/api/v1/user/keys' % get_url('Gogs', None)
    username, password = get_username_password('Gogs')
    auth = HTTPBasicAuth(username, password)
    response = get(keys_url, auth=auth)
    if response.status_code != 200:
        raise bootstrap_error(response)
    keys = response.json()
    local_key = get_local_key()
    if not key_exists(keys, local_key):
        response = post(
            keys_url, auth=auth, json={'title': node(),
                                       'key': local_key})
        if response.status_code != 201:
            raise bootstrap_error(response)


def set_ssh_keys():
    """Interactively set SSH keys on remote Git servers."""
    try:
        for service in ['GitHub', 'GitLab', 'BitBucket Cloud', 'Gogs']:
            question = 'Set %s SSH key' % service
            default = 'Y'
            while agree(question, default):
                try:
                    {
                        'GitHub': set_github_ssh_key,
                        'GitLab': set_gitlab_ssh_key,
                        'BitBucket Cloud': set_bitbucket_cloud_ssh_key,
                        'Gogs': set_gogs_ssh_key,
                    }[service]()
                    if service == 'BitBucket Cloud':
                        break
                    question = 'Set another %s SSH key' % service
                    default = 'N'
                except (BootstrapError, ConnectionError) as error:
                    print('error: %s' % error.message)
    except KeyboardInterrupt:
        exit(130)