pypump-shell

#!/usr/bin/env python
##
# Copyright (C) 2013 Jessica T. (Tsyesika) <xray7224@googlemail.com>
#
# This program 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/>.
##

import logging
import argparse
import sys

import requests
import six

from pypump import PyPump, Client

try:
    from IPython import embed
    from IPython.config.loader import Config
except ImportError:
    six.print_("You need to have ipython installed to run this")
    sys.exit(1)

welcome_banner = """Welcome to the PyPump Shell,
We have setup some useful objects for you:

{table}

If you need help please visit our docs:
https://pypump.readthedocs.org/en/latest
"""


def ascii_table(headings, data):
    """ Draws an ascii table """
    # first we need to work out how long each column should be
    columns = {}
    for heading in headings:
        columns[heading] = len(heading)+1  # plus one for the extra padding

    # data could also overspill
    for record in data:
        for heading, d in record.items():
            dlen = len(d)+1
            if dlen > columns[heading]:
                columns[heading] = dlen

    table = {}
    # okay now we need to pad the headers readying for drawing
    for column, width in columns.items():
        table[column] = "{name}{padding}".format(
            name=column, padding=" "*(width-len(column))
            )

    heading = "| "
    for column in table.values():
        heading += column
        heading += "| "

    # now for the data
    table = {}
    for i, record in enumerate(data):
        table[i] = "| "
        for column, cdata in record.items():
            table[i] += "{data}{padding} | ".format(
                data=cdata, padding=" "*(columns[column]-len(cdata)-1)
                )

    # make the helper function which will draw the seporators
    def draw_sep(columns):
        """ Draws +----+--- etc... """
        sep = "+"
        for width in columns:
            sep += "-"*(width+1)
            sep += "+"

        return sep + "\r\n"

    sepper = lambda: draw_sep(columns.values())
    stable = sepper()
    stable += heading + "\r\n"
    stable += sepper()
    for value in table.values():
        stable += value + "\r\n"

    stable += sepper()
    return stable  # few, glad that's over


def verifier(url):
    """ Asks for verification code for OAuth OOB """
    six.print_("Please open and follow the instructions:")
    six.print_(url)
    return six.moves.input("Verifier: ").lstrip(" ").rstrip(" ")

if __name__ == "__main__":
    # Taken from https://docs.python.org/2/library/logging.html#logging-levels
    log_levels =  "{'CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'}"

    parser = argparse.ArgumentParser(description='Command line interface to the pump.io APIs')

    parser.add_argument("--log", "--loglevel",
                        dest="loglevel",
                        default="warning",
                        help="Set level of logging. Can be one of {}".format(log_levels))

    parser.add_argument("--logfile",
                        dest="logfile",
                        action="store_true",
                        default=False,
                        help="Write logging to pypump-shell.log")

    parser.add_argument("--no-check-certificate",
                        dest="verify",
                        action="store_false",
                        default=True,
                        help="Accepts invalid SSL certificates")

    parser.add_argument("webfinger",
                        help="Webfinger to use when connecting")

    options = parser.parse_args()

    log_level = getattr(logging, options.loglevel.upper(), None)
    if not isinstance(log_level, int):
        raise ValueError('Invalid log level: {}. Can be one of {}'.format(options.loglevel,
                                                                          log_levels))

    logfile = "pypump-shell.log" if options.logfile else None
    logging.basicConfig(level=log_level, filename=logfile)

    webfinger = options.webfinger

    client = Client(
        webfinger=webfinger,
        name="PyPump Shell",
        type="native"
    )

    sys.stdout.write("-> Setting up PyPump  ")
    sys.stdout.flush()  # actually get it on the screen - most terms wait for \n

    try:
        pump = PyPump(client=client, verifier_callback=verifier, verify_requests=options.verify)
    except requests.exceptions.SSLError:
        # This is caused when there has been an invalid certificate.
        # We should ask the user if they want to continue (common to use
        # self-signed/invalid certificates in dev enviroment).
        six.print_("Couldn't validate SSL/TLS certificate for {0}".format(webfinger.split("@", 1)[1]))
        answer = None
        while answer not in ["y", "n"]:
            if answer is not None:
                six.print_("Sorry, please type 'y' or 'n' (Ctrl-C to exit).")

            answer = raw_input("Would you like to accept the invalid/untrusted certificate (y/n): ").lower()
            answer = answer.replace(" ", "")

        if answer == "n":
            sys.exit(0)

        # Create the same PyPump object as above but without ssl verification
        pump = PyPump(client=client, verifier_callback=verifier, verify_requests=False)

    # bring curser back so banner walks over the setup message
    sys.stdout.write("\r")

    # drop them into a shell
    cfg = Config()
    cfg.InteractiveShell.confirm_exit = False

    # prep the welcome banner
    welcome_banner = welcome_banner.format(
        table=ascii_table(
            ["Variable", "Representation", "Type"],
            [
                {
                    "Variable": "pump",
                    "Representation": str(repr(pump)),
                    "Type": type(pump).__name__,
                },
            ]
        )
    )

    embed(
        config=cfg,
        banner1=welcome_banner,
        exit_msg="Remeber! Report any bugs to https://github.com/xray7224/PyPump/issues"
    )