jupyter_o2

#!/usr/bin/env python
# author: Aaron Kollasch
# date: 2018-2-8
"""A script to launch and connect to a Jupyter session on Orchestra 2.
as described at https://wiki.rc.hms.harvard.edu/display/O2/Jupyter+on+O2

Usage: jupyter_remote <USER> <subcommand>

Example: jupyter_remote js123 notebook
    This will launch an X11-enabled ssh, start an interactive node running jupyter notebook,
    ssh into that interactive node to allow requests to be forwarded,
    and finally open the notebook in your browser.

Troubleshooting:
    If jupyter hangs when opening notebooks for the first time in any session,
    and the console shows a "The signatures database cannot be opened" error,
    enter an interactive session and generate a notebook config using
      `jupyter notebook --generate-config`
    then set c.NotebookNotary.db_file = ':memory:'

    If you see `srun: error: x11: no local DISPLAY defined`,
    try logging in to the server with `ssh -X` and check your DISPLAY using `echo $DISPLAY`.
    There should be a string printed in response.
    Otherwise, try reinstalling XQuartz at https://www.xquartz.org/
    or run jupyter_remote with the -Y argument to enable trusted X11 forwarding.
"""

from __future__ import print_function
import sys
import os
import argparse
import getpass  # fallback if pinentry is not installed (optional "brew install pinentry" on macs)
import subprocess
import re
import webbrowser
import shlex
import ctypes
import atexit
from signal import signal, SIGABRT, SIGINT, SIGTERM
try:
    from ConfigParser import SafeConfigParser as ConfigParser
except ImportError:
    from configparser import ConfigParser
try:
    from shlex import quote
except ImportError:
    from pipes import quote

from pexpect import pxssh
try:
    # optional import of dnspython to resolve "Unknown host" errors
    # if you get errors like "Could not establish connection to host",
    # tests the server with nslookup <host>
    # your current DNS server may not have entries for the individual O2 login nodes
    import dns.resolver
except ImportError:
    dns = None

#############################################################
# program-wide and system-specific defaults                 #
#############################################################

DEFAULT_CFG = """[Defaults]
# Uncomment lines in this section to override the default options 
# by deleting the corresponding semicolons:
;DEFAULT_HOST = o2.hms.harvard.edu
;DEFAULT_JP_PORT = 8887
;DEFAULT_JP_TIME = 0-12:00
;DEFAULT_JP_MEM = 1G
;DEFAULT_JP_CORES = 1
;DEFAULT_JP_SUBCOMMAND = notebook

[Settings]
MODULE_LOAD_CALL = gcc/6.2.0 python/2.7.12
# The command to load the version of python used by your virtual environment.
# This setting will direct jupyter_remote to send `module load <MODULE_LOAD_CALL>`
# before sending `source <SOURCE_JUPYTER_CALL>`
# Subtract out `gcc/6.2.0 python/2.7.12` if MODULE_LOAD_CALL is not needed.

SOURCE_JUPYTER_CALL = jupytervenv/bin/activate
# The bash source command to enter your jupyter environment.
# Jupyter will send `source <SOURCE_JUPYTER_CALL`, so adding source is not necessary
# SOURCE_JUPYTER_CALL will be single quoted to block injection. 
# Thus, it does not support tilde expansion.
# (jupyter_remote should start in the home directory so this shouldn't be necessary)
"""

JO2_DEFAULTS = {
    "DEFAULT_HOST": "o2.hms.harvard.edu",
    "DEFAULT_JP_PORT": "8887",
    "DEFAULT_JP_TIME": "0-12:00",
    "DEFAULT_JP_MEM": "1G",
    "DEFAULT_JP_CORES": "1",
    "DEFAULT_JP_SUBCOMMAND": "notebook",
    "MODULE_LOAD_CALL": "",
    "SOURCE_JUPYTER_CALL": "",
}

_config = ConfigParser(defaults=JO2_DEFAULTS)
_config.add_section('Defaults')
_config.add_section('Settings')

CFG_FILENAME = ".jupyter-remote.cfg"
CFG_FULLPATH = os.path.join(os.path.expanduser("~"), CFG_FILENAME)

if sys.version_info.major >= 3:
    try:
        with open(CFG_FULLPATH, 'x') as cfg_file:
            print("Configuration file not found at {}".format(CFG_FULLPATH))
            print("Creating {}".format(CFG_FILENAME))
            cfg_file.write(DEFAULT_CFG)
    except (IOError, FileExistsError):
        pass
else:
    # Check if the .cfg exists before writing
    # Not using EAFP because there are many reasons why config.read() would return []
    if not os.path.isfile(CFG_FULLPATH):
        print("Configuration file not found at {}".format(CFG_FULLPATH))
        print("Creating {}".format(CFG_FILENAME))
        try:
            with open(CFG_FULLPATH, 'w') as cfg_file:
                cfg_file.write(DEFAULT_CFG)
        except IOError:
            pass

if not _config.read(CFG_FULLPATH):
    print("Config file could not be read. Using internal defaults.")

DEFAULT_HOST = _config.get('Defaults', 'DEFAULT_HOST')
DEFAULT_JP_PORT = _config.getint('Defaults', 'DEFAULT_JP_PORT')
DEFAULT_JP_TIME = _config.get('Defaults', 'DEFAULT_JP_TIME')
DEFAULT_JP_MEM = _config.get('Defaults', 'DEFAULT_JP_MEM')
DEFAULT_JP_CORES = _config.getint('Defaults', 'DEFAULT_JP_CORES')
DEFAULT_JP_SUBCOMMAND = _config.get('Defaults', 'DEFAULT_JP_SUBCOMMAND')

MODULE_LOAD_CALL = _config.get('Settings', 'MODULE_LOAD_CALL')
SOURCE_JUPYTER_CALL = _config.get('Settings', 'SOURCE_JUPYTER_CALL')

JO2_ARG_PARSER = argparse.ArgumentParser(description='Launch and connect to a Jupyter session on O2.')
JO2_ARG_PARSER.add_argument("user", type=str, help="O2 username")
JO2_ARG_PARSER.add_argument("subcommand", help="the subcommand to launch")
JO2_ARG_PARSER.add_argument("--host", type=str, default=DEFAULT_HOST, help="Host to connect to")
JO2_ARG_PARSER.add_argument("-p", "--port", dest="jp_port", type=int, default=DEFAULT_JP_PORT,
                            help="Available port on your system")
JO2_ARG_PARSER.add_argument("-t", "--time", dest="jp_time", type=str, default=DEFAULT_JP_TIME,
                            help="Time to run Jupyter")
JO2_ARG_PARSER.add_argument("-m", "--mem", dest="jp_mem", type=str, default=DEFAULT_JP_MEM,
                            help="Memory to allocate for Jupyter")
JO2_ARG_PARSER.add_argument("-c", "-n", dest="jp_cores", type=int, default=DEFAULT_JP_CORES,
                            help="Cores to allocate for Jupyter")
JO2_ARG_PARSER.add_argument("-k", "--keepalive", default=False, action='store_true',
                            help="Keep interactive session alive after exiting Jupyter")
JO2_ARG_PARSER.add_argument("--kq", "--keepxquartz", dest="keepxquartz", default=False, action='store_true',
                            help="Do not quit XQuartz")
JO2_ARG_PARSER.add_argument("-Y", "--ForwardX11Trusted", dest="forwardx11trusted", default=False, action='store_true',
                            help="Enables trusted X11 forwarding. Equivalent to ssh -Y.")

SRUN_CALL_FORMAT = "srun -t {} --mem {} -c {} --pty -p interactive --x11 /bin/bash"
JP_CALL_FORMAT = "jupyter {} --port={} --browser='none'"

# patterns to search for in ssh
SITE_PATTERN_FORMAT = "\s(https?://((localhost)|(127\.0\.0\.1)):{}[\w\-./%?=]+)\s"  # {} formatted with jupyter port
PASSWORD_PATTERN = re.compile(b"[\w-]+@[\w-]+'s password: ")  # e.g. "user@compute-e-16-175's password: "

if sys.platform.startswith("linux"):
    PINENTRY_PATH = "/usr/bin/pinentry"
elif sys.platform == "darwin":
    PINENTRY_PATH = "/usr/local/bin/pinentry"
else:
    PINENTRY_PATH = "pinentry"

if dns is not None:
    DNS_SERVER_GROUPS = [  # dns servers that have entries for loginXX.o2.rc.hms.harvard.edu
        dns.resolver.Resolver().nameservers,                    # current nameservers, checked first
        ["134.174.17.6", "134.174.141.2"],                      # HMS nameservers
        ["128.103.1.1", "128.103.201.100", "128.103.200.101"],  # HU nameservers
    ]
    # tests that you can access the login nodes with nslookup login01.o2.rc.hms.harvard.edu <DNS>
else:
    DNS_SERVER_GROUPS = None

if sys.version_info.major >= 3:
    STDOUT_BUFFER = sys.stdout.buffer
else:
    STDOUT_BUFFER = sys.stdout


#####################################################
# functions and classes                             #
#####################################################


class JupyterO2(object):
    def __init__(
            self,
            user,
            host,
            subcommand=DEFAULT_JP_SUBCOMMAND,
            jp_port=DEFAULT_JP_PORT,
            jp_time=DEFAULT_JP_TIME,
            jp_mem=DEFAULT_JP_MEM,
            jp_cores=DEFAULT_JP_CORES,
            keepalive=False,
            keepxquartz=False,
            forwardx11trusted=False,
    ):
        self.user = user
        self.host = host
        self.subcommand = subcommand
        self.jp_port = jp_port
        self.keep_alive = keepalive
        self.keep_xquartz = keepxquartz

        self.srun_call = SRUN_CALL_FORMAT.format(quote(jp_time), quote(jp_mem), jp_cores)
        self.jp_call = JP_CALL_FORMAT.format(quote(subcommand), jp_port)

        self.__o2_pass = ""
        self.__pinentry = Pinentry(pinentry_path=PINENTRY_PATH, fallback_to_getpass=True)

        login_ssh_options = {
            "ForwardX11": "yes",
            "LocalForward": "{} 127.0.0.1:{}".format(jp_port, jp_port),
            "PubkeyAuthentication": "no"
        }
        if forwardx11trusted:
            login_ssh_options["ForwardX11Trusted"] = "yes"

        self.__login_ssh = CustomSSH(timeout=60, ignore_sighup=False, options=login_ssh_options)

        self.__second_ssh = CustomSSH(timeout=10, ignore_sighup=False, options={"PubkeyAuthentication": "no"})

        # perform close() on exit or interrupt
        atexit.register(self.close)
        self.flag_exit = False
        for sig in (SIGABRT, SIGINT, SIGTERM):
            signal(sig, self.term)

    def ask_for_pin(self):
        self.__o2_pass = self.__pinentry.ask(
            prompt="Enter your passphrase: ",
            description="Connect to O2 server for jupyter {}".format(self.subcommand),
            error="No password entered", validator=lambda x: x is not None and len(x) > 0)
        self.__pinentry.close()

    def connect(self):
        """
        First SSH into an interactive node and run jupyter.
        Then SSH into that node to set up forwarding.
        Finally, open the jupyter notebook page in the browser.
        """
        # start login ssh
        print("ssh {}@{}".format(self.user, self.host))
        dns_err, host = check_dns(self.host)
        if dns_err == 1:
            print("ssh {}@{}".format(self.user, host))
        elif dns_err == 2:
            sys.exit(1)
        self.__login_ssh.force_password = True
        self.__login_ssh.silence_logs()
        self.__login_ssh.login(host, self.user, self.__o2_pass)

        # get the login hostname
        self.__login_ssh.sendline("hostname")
        self.__login_ssh.prompt()
        jp_login_host = self.__login_ssh.before.decode('utf-8').strip().split('\n')[1]
        print("hostname: {}".format(jp_login_host))

        # enter an interactive session
        print()
        self.__login_ssh.PROMPT = PASSWORD_PATTERN
        self.__login_ssh.logfile_read = FilteredOut(STDOUT_BUFFER, b'srun')
        self.__login_ssh.sendline(self.srun_call)
        if not self.__login_ssh.prompt():
            eprint("The timeout ({}) was reached.".format(self.__login_ssh.timeout))
            sys.exit(1)
        self.__login_ssh.silence_logs()
        self.__login_ssh.sendline(self.__o2_pass)

        # within interactive: get the name of the interactive node
        self.__login_ssh.PROMPT = self.__login_ssh.UNIQUE_PROMPT
        self.__login_ssh.sendline("unset PROMPT_COMMAND; PS1='[PEXPECT]\$ '")
        self.__login_ssh.prompt()
        self.__login_ssh.sendline("hostname | sed 's/\..*//'")
        self.__login_ssh.prompt()
        jp_interactive_host = self.__login_ssh.before.decode('utf-8').strip().split('\n')[1]
        print("interactive node: {}\n".format(jp_interactive_host))

        # start jupyter
        if MODULE_LOAD_CALL:
            print(join_cmd("module load", MODULE_LOAD_CALL))
            self.__login_ssh.sendline(join_cmd("module load", MODULE_LOAD_CALL))
            self.__login_ssh.prompt()
        if SOURCE_JUPYTER_CALL:
            print(join_cmd("source", SOURCE_JUPYTER_CALL))
            self.__login_ssh.sendline(join_cmd("source", SOURCE_JUPYTER_CALL))
            self.__login_ssh.prompt()
        self.__login_ssh.sendline(self.jp_call)
        self.__login_ssh.logfile_read = STDOUT_BUFFER

        # get the address jupyter is running at
        site_pat = re.compile(SITE_PATTERN_FORMAT.format(self.jp_port).encode('utf-8'))
        self.__login_ssh.PROMPT = site_pat
        if not self.__login_ssh.prompt():  # timed out; failed to launch jupyter
            eprint("Failed to launch jupyter. (timed out, {})".format(self.__login_ssh.timeout))
            if self.keep_alive:
                eprint("Starting interactive mode.")
                self.interact()
            else:
                sys.exit(1)
        jp_site = self.__login_ssh.after.decode('utf-8').strip()

        # log in to the second ssh
        print("\nssh {}@{}".format(self.user, jp_login_host))
        dns_err, jp_login_host = check_dns(jp_login_host)
        if dns_err == 1:
            print("ssh {}@{}".format(self.user, jp_login_host))
        elif dns_err == 2:
            sys.exit(1)
        self.__second_ssh.force_password = True
        self.__second_ssh.silence_logs()
        self.__second_ssh.login(jp_login_host, self.user, self.__o2_pass)

        # ssh into the running interactive node
        print("\nssh -N -L {0}:127.0.0.1:{0} {1}".format(self.jp_port, jp_interactive_host))
        self.__second_ssh.PROMPT = PASSWORD_PATTERN
        self.__second_ssh.sendline("ssh -N -L {0}:127.0.0.1:{0} {1}".format(self.jp_port, jp_interactive_host))
        if not self.__second_ssh.prompt():
            eprint("The timeout ({}) was reached.".format(self.__second_ssh.timeout))
            sys.exit(1)
        self.__second_ssh.silence_logs()
        self.__second_ssh.sendline(self.__o2_pass)
        zero(self.__o2_pass)  # password is not needed anymore
        self.__second_ssh.logfile_read = STDOUT_BUFFER  # print any errors/output from self.__second_ssh to stdout

        # open Jupyter in browser
        print("\nJupyter is ready! Access at:\n{}\nOpening in browser...\n".format(jp_site))
        try:
            webbrowser.open(jp_site, new=2)
        except webbrowser.Error as error:
            print("Error: {}\nPlease open the Jupyter page manually.".format(error))

        # quit XQuartz because the application is not necessary to keep the connection open.
        if not self.keep_xquartz:
            try_quit_xquartz()

    def interact(self):
        """Keep the ssh session alive and allow input such as Ctrl-C to close Jupyter."""
        self.__login_ssh.silence_logs()
        if self.keep_alive:  # exit when you log out of the login shell
            interact_filter = FilteredOut(None, b'[PEXPECT]$ logout')
            self.__login_ssh.interact(output_filter=interact_filter.exit_on_find)
        else:  # exit when jupyter exits and [PEXPECT]$ appears
            interact_filter = FilteredOut(None, b'[PEXPECT]$ ')
            self.__login_ssh.interact(output_filter=interact_filter.exit_on_find)

    def close(self, cprint=print, *_):
        """cprint allows printing to be disabled if necessary using `cprint=lambda x, end=None, flush=None: None`"""
        def _cprint(*args, **kwargs):
            if sys.version_info.major == 2:
                kwargs.pop('flush', None)
            cprint(*args, **kwargs)
        _cprint("Cleaning up\r\n", end="", flush=True)
        zero(self.__o2_pass)
        self.__pinentry.close()
        if not self.__login_ssh.closed:
            _cprint("Closing login_ssh\n", end="", flush=True)
            self.__login_ssh.close(force=True)
        if not self.__second_ssh.closed:
            _cprint("Closing second_ssh\n", end="", flush=True)
            self.__second_ssh.close(force=True)

    def term(self, *_):
        if not self.flag_exit:
            self.flag_exit = True
            try:
                self.close()
            except RuntimeError:  # printing from signal can cause RuntimeError: reentrant call
                self.close(cprint=lambda x, end=None, flush=None: None)
            sys.stdout.close()
            sys.stderr.close()
            sys.stdin.close()
            os.closerange(0, 3)
            os._exit(1)


def join_cmd(cmd, args_string):
    """Create a bash command by joining cmd and args_string. Resistant to injection within args_string."""
    return ' '.join([cmd] + [quote(item) for item in shlex.split(args_string)])


def check_dns(hostname, dns_groups=DNS_SERVER_GROUPS):
    """Check if hostname is reachable by any group of dns servers."""
    if dns is not None:
        dns_err_code = 0
        for dns_servers in dns_groups:
            try:
                my_resolver = dns.resolver.Resolver()
                my_resolver.nameservers = dns_servers
                if dns_err_code > 0:
                    eprint("Could not resolve domain. Trying with nameservers: {}".format(dns_servers))
                    answer = my_resolver.query(hostname)
                    hostname = answer[0].address
                    dns_err_code = 1
                else:
                    my_resolver.query(hostname)
                break
            except dns.resolver.NXDOMAIN:
                dns_err_code = 2
    else:
        dns_err_code = -1
    if dns_err_code == 1:
        print("Found IP: {}".format(hostname))
    elif dns_err_code == 2:
        eprint("No IP found for {}".format(hostname))
    return dns_err_code, hostname


class CustomSSH(pxssh.pxssh):
    def login(self, *args, **kwargs):
        """Login and suppress the traceback for pxssh exceptions (such as incorrect password errors)."""
        try:
            super(CustomSSH, self).login(*args, **kwargs)
        except pxssh.ExceptionPxssh as err:
            eprint("pxssh error: {}".format(err))
            sys.exit(1)

    def silence_logs(self):
        """Prevent printing into any logfile."""
        self.logfile = None
        self.logfile_read = None
        self.logfile_send = None

    def digest_all_prompts(self, timeout=0.5):
        """Digest all prompts until there is a delay of <timeout>."""
        if timeout == -1:
            timeout = self.timeout
        while self.prompt(timeout):
            pass


class FilteredOut(object):
    def __init__(self, txtctrl, by):
        self.txtctrl = txtctrl
        self.by = by

    def write(self, bytestr):
        try:
            if bytestr[:len(self.by)] == self.by:
                self.txtctrl.write(bytestr)
        except IndexError:
            pass

    def flush(self):
        self.txtctrl.flush()

    def exit_on_find(self, bytestr):
        if self.by in bytestr:
            sys.exit(0)
        return bytestr


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def try_quit_xquartz():
    """Quit XQuartz on macs
    First attempts to check if there is a window open in XQuartz, using Quartz (pyobjc-framework-Quartz).
    If Quartz is installed, and there is a window open in XQuartz, it will not quit XQuartz.
    If Quartz is not installed and there is a window open in XQuartz,
        XQuartz open a dialog to ask if you really want to quit it.
    Otherwise XQuartz will silently quit.
    """
    if sys.platform != "darwin":
        return

    try:
        if not xquartz_is_open():
            return
        print("Quitting XQuartz... ", end='')
        open_windows = get_xquartz_open_windows()
        if open_windows is None:
            pass
        elif not open_windows:
            quit_xquartz()
        else:
            print("\nXQuartz window(s) are open. Not quitting.")
            try:
                print("Open windows: {}".format(
                    [(window['kCGWindowName'], window['kCGWindowNumber']) for window in open_windows]))
            except KeyError:
                pass
        if not xquartz_is_open():
            print("Success.")
        else:
            print("Failed to quit.")
    except Exception as error:
        eprint("Error: {}".format(error.__class__))
        eprint(error)
        print("Failed to quit XQuartz.")


def get_xquartz_open_windows():
    """
    Get info on all open XQuartz windows.
    Requires pyobjc-framework-Quartz (install with pip)
    :return: a list of open windows as python dictionaries
    """
    try:
        from Quartz import CGWindowListCopyWindowInfo, kCGWindowListOptionAll, kCGNullWindowID
        from PyObjCTools import Conversion

        # need to use kCGWindowListOptionAll to include windows that are not currently on screen (e.g. minimized)
        windows = CGWindowListCopyWindowInfo(kCGWindowListOptionAll, kCGNullWindowID)

        # then filter for XQuartz main windows
        open_windows = [window for window in windows if window['kCGWindowOwnerName'] == "XQuartz" and
                        window['kCGWindowLayer'] == 0 and 'kCGWindowName' in window.keys() and
                        window['kCGWindowName'] not in ['', 'X11 Application Menu', 'X11 Preferences']]

        # convert from NSDictionary to python dictionary
        open_windows = Conversion.pythonCollectionFromPropertyList(open_windows)

        return open_windows
    except (ImportError, KeyError) as error:
        eprint("Error: {}".format(error.__class__))
        eprint(error)
        return None


def xquartz_is_open():
    try:
        proc = subprocess.check_output(["pgrep", "-f", "XQuartz"])
        return len(proc) > 0
    except subprocess.CalledProcessError as error:
        if error.returncode == 1:
            return False  # pgrep returns with exit code 1 if XQuartz is not open
        else:  # syntax error or fatal error
            raise error


def quit_xquartz():
    if sys.platform == "darwin":
        subprocess.call(['osascript', '-e', 'quit app "XQuartz"'])


#################################################################
# pin entry and security functions are slightly modified        #
# from pysectools (Greg V <greg@unrelenting.technology>):       #
#                                                               #
# - updated pysectools for python 3:                            #
#     uses bytestrings and flushes stdin after writing          #
# - removed shell=True from subprocess calls                    #
#################################################################


def zero(s):
    """
    Tries to securely erase a secret string from memory
    (overwrite it with zeros.)

    Only works on CPython.

    Returns True if successful, False otherwise.
    """
    try:
        bufsize = len(s) + 1
        offset = sys.getsizeof(s) - bufsize
        location = id(s) + offset
        ctypes.memset(location, 0, bufsize)
        return True
    except Exception:
        return False


def cmd_exists(cmd):
    return subprocess.call(["type", cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0


class PinentryException(Exception):
    pass


class PinentryUnavailableException(PinentryException):
    pass


class PinentryClosedException(PinentryException):
    pass


class PinentryErrorException(PinentryException):
    pass


class Pinentry(object):
    def __init__(self, pinentry_path="pinentry", fallback_to_getpass=True):
        if not cmd_exists(pinentry_path):
            if fallback_to_getpass and os.isatty(sys.stdout.fileno()):
                self._ask = self._ask_with_getpass
                self._close = self._close_getpass
            else:
                raise PinentryUnavailableException()
        else:
            self.process = subprocess.Popen(pinentry_path,
                                            stdin=subprocess.PIPE,
                                            stdout=subprocess.PIPE,
                                            stderr=subprocess.STDOUT,
                                            close_fds=True)
            self._ask = self._ask_with_pinentry
            self._close = self._close_pinentry
        self._closed = False

    def ask(self,
            prompt="Enter the password: ",
            description=None,
            error="Wrong password!",
            validator=lambda x: x is not None):
        if self._closed:
            raise PinentryClosedException()
        return self._ask(prompt, description, error, validator)

    def close(self):
        self._closed = True
        return self._close()

    @staticmethod
    def _ask_with_getpass(prompt, description, error, validator):
        if description:
            print(description, flush=True)
        password = None
        while not validator(password):
            if password is not None:
                eprint(error)
            password = getpass.getpass(prompt)
        return password

    def _close_getpass(self): pass

    def _ask_with_pinentry(self, prompt, description, error, validator):
        self._waitfor("OK")
        env = os.environ.get
        self._comm("OPTION lc-ctype=%s" % env("LC_CTYPE", env("LC_ALL", "en_US.UTF-8")))
        try:
            self._comm("OPTION ttyname=%s" % env("TTY", os.ttyname(sys.stdout.fileno())))
        except Exception:
            pass
        if env('TERM'):
            self._comm("OPTION ttytype=%s" % env("TERM"))
        if prompt:
            self._comm("SETPROMPT %s" % self._esc(prompt))
        if description:
            self._comm("SETDESC %s" % self._esc(description))
        password = None
        while not validator(password):
            if password is not None:
                self._comm("SETERROR %s" % self._esc(error))
            self.process.stdin.write(b"GETPIN\n")
            self.process.stdin.flush()
            try:
                password = self._waitfor("D ", breakat="OK", errat="ERR")
            except PinentryErrorException:
                sys.exit(0)
            if password is not None:
                password = password[2:].replace("\n", "")
        return password

    def _close_pinentry(self):
        return self.process.kill()

    def _waitfor(self, what, breakat=None, errat=None):
        out = ""
        while not out.startswith(what):
            if breakat is not None and out.startswith(breakat):
                break
            elif errat is not None and out.startswith(errat):
                raise PinentryErrorException()
            out = self.process.stdout.readline().decode('utf-8')
        return out

    def _comm(self, x):
        self.process.stdin.write(x.encode('utf-8') + b"\n")
        self.process.stdin.flush()
        self._waitfor("OK")

    @staticmethod
    def _esc(x):
        return x.replace("%", "%25").replace("\n", "%0A")


#############################################################
# run jupyter_remote if it is the main program                  #
#############################################################

if __name__ == "__main__":
    pargs = JO2_ARG_PARSER.parse_args()

    jupyter_o2 = JupyterO2(**vars(pargs))

    jupyter_o2.ask_for_pin()
    jupyter_o2.connect()
    jupyter_o2.interact()