start_jupyter

"""
This script handles four different modes.

1. start_jupyter
  Runs notebook server in current directory

2. start_jupyter foo.ipynb
  Runs notebook in current directory

3. /path/to/tool/middleware/invoke
  Runs tools.  The invoke script will give the
  path to the notebook to run.  The path of the
  invoke script is used to locate the notebook.
  Current directory will be set to the notebook
  location (usually in the tool bin directory)

4. jupyter/middleware/invoke
  Run jupyter tool.  Starts a notebook server in
  your "notebooks" directory.
"""

from __future__ import print_function
import sys
import os
import re
import glob
import shutil
import signal
import argparse
import json
import binascii
from subprocess import call, Popen, check_output
from string import Template
import time

notebookProcess = None

def handler(signum, frame):
    global notebookProcess

    print("%s Signal %d received.\n" % (time.ctime(),signum))
    sys.stdout.flush()  # make sure this gets displayed now.
    if notebookProcess:
        try:
            print("%s killing notebookProcess = %d" % (time.ctime(),notebookProcess.pid))
            sys.stdout.flush()  # make sure this gets displayed now.
#           os.kill(notebookProcess.pid,signal.SIGTERM)
            notebookProcess.terminate()
        except:
            print("%s killing notebookProcess = %d failed" % (time.ctime(),notebookProcess.pid))
            sys.stdout.flush()  # make sure this gets displayed now.


def completeHandler(signum, frame):
    global notebookProcess

    print("%s child process killed(%d)" % (time.ctime(),signum))
    sys.stdout.flush()  # make sure this gets displayed now.


signal.signal(signal.SIGINT, handler)
signal.signal(signal.SIGHUP, handler)
signal.signal(signal.SIGTERM, handler)
signal.signal(signal.SIGCHLD, completeHandler)


def parse_cmd_line():
    prog = "start_jupyter"

    parser = argparse.ArgumentParser(
        usage="""usage: %(prog)s [-h] [-t] [-c] [-d] [-A] [-T dir] --themes [name]

Start a Jupyter notebook-based tool

positional arguments:
  name        Name of notebook to run.  The terminal and
              dashboard (tree view of files) will de disabled.

              If no name is given, a notebook server will
              be started in the current directory.  Terminal
              and dashboard will be enabled.

optional arguments:
  -h, --help  show this help message and exit.
  -d          Show debug (verbose) output.
  -t          Run as a Tool with no notebook controls.
  -c          Copy instead of link notebook files.
  -A          Run in AppMode.
  -T dir      Search for notebook starting in dir.
  --themes    Enable notebook themes
""",
        prog=prog,
        add_help=False)
    parser.add_argument('-t', dest='tool', action='store_true')
    parser.add_argument('-A', dest='app', action='store_true')
    parser.add_argument('-d', dest='debug', action='store_true')
    parser.add_argument('-c', dest='copyNotebooks', action='store_true')
    parser.add_argument('-T')
    parser.add_argument('-h', '--help', dest='help', action='store_true')
    parser.add_argument('name', nargs='?')
    return parser


def get_session():
    try:
        session = os.environ['SESSION']
        sessiondir = os.environ['SESSIONDIR']
    except:
        session = ""

    if session != "":
        return session, sessiondir

    # something went wrong. restore environment variables
    psout = check_output(['ps', 'auxwwe']).split('\n')
    for line in psout:
        ind = line.find('SESSIONDIR')
        if ind > 0:
            sessiondir = line[ind:].split(' ')[0]
            display = line[line.find('DISPLAY'):].split(' ')[0]
            break
    session = sessiondir.split('/')[-1]
    sessiondir = sessiondir.split('=')[1]
    display = display.split('=')[1]
    os.environ['SESSION'] = session
    os.environ['SESSIONDIR'] = sessiondir
    os.environ['DISPLAY'] = display
    return session, sessiondir

def get_proxy_addr():
    session, sessiondir = get_session()
    fn = os.path.join(sessiondir, 'resources')
    with open(fn, 'r') as f:
        res = f.read()
    for line in res.split('\n'):
        if line.startswith('hub_url'):
            url = line.split()[1]
        elif line.startswith('filexfer_port'):
            fxp = str(int(line.split()[1]) % 1000)
        elif line.startswith('filexfer_cookie'):
            fxc = line.split()[1]
    url_path = "/weber/%s/%s/%s/" % (session, fxc, fxp)
    proxy_url = "https://proxy." + url.split('//')[1] + url_path
    return url_path, proxy_url


def get_cookie():
    cookie_name = ""
    try:
        session = int(os.environ['SESSION'])
        pwfile = glob.glob('/var/run/Xvnc/passwd-*')[0]
        with open(pwfile, 'rb') as f:
            pwd = binascii.hexlify(f.read()).decode('utf-8')
            token = "%d:%s" % (session, str(pwd))

        fn = os.path.join(os.environ['SESSIONDIR'], 'resources')
        with open(fn, 'r') as f:
            res = f.read()
        for line in res.split('\n'):
            if line.startswith('hub_url'):
                url = line.split()[1]
                host = url[url.find('//') + 2:]
                cookie_name = 'weber-auth-' + host.replace('.', '-')
                break
    except:
        # not running on a hub
        return "", ""
    return cookie_name, token


def find_notebook(name, tool, copyNotebooks, arglist):
    if tool is None:
        if not os.path.isfile(name):
            print("Cannot find %s." % name, file=sys.stderr)
            sys.exit(1)
        return os.path.split(name)

    # We are running a published notebook.
    # symlink copy it to RESULTSDIR
    tname = os.path.basename(os.path.dirname(tool))
    tdest = os.path.join(os.environ['RESULTSDIR'], tname)

    # remove and make tdest
    if os.path.exists(tdest):
        shutil.rmtree(tdest)
    os.makedirs(tdest)

    realPath = os.path.realpath(os.path.abspath(tool))
    call("cp -sRf %s/* %s" % (realPath,tdest), shell=True)
    if copyNotebooks:
        call("""( cd %s; find ./ -name "*.ipynb" -exec /bin/cp --remove-destination {} %s/{} \; )""" % (realPath,tdest),shell=True)

    # notebook will be either relative to the tool
    # directory or in the bin directory
    nb_name = os.path.abspath(os.path.join(tdest, name))
    if not os.path.isfile(nb_name):
        # look in bin dir
        nb_name = os.path.abspath(os.path.join(tdest, 'bin', name))
        if not os.path.isfile(nb_name):
            print("Cannot find %s." % name, file=sys.stderr)
            sys.exit(1)

    return os.path.split(nb_name)


def run_notebook(args):
    global notebookProcess

    arglist = []

    if args.name is None:
        if args.tool:
            parser.error("ERROR: -t with no filename.\n")
        if args.T:
            os.chdir(os.environ['HOME'])

    else:
        nb_dir, nb_name = find_notebook(args.name, args.T, args.copyNotebooks, arglist)
        if nb_dir:
            print("Changing directory to", nb_dir)
            os.chdir(nb_dir)

        if args.app:
            arglist.append('--NotebookApp.default_url="apps/%s"' % nb_name)
        else:
            arglist.append('--NotebookApp.default_url="notebooks/%s"' % nb_name)
# disable ~/.local/lib/pythonX.Y/site-packages for tools
        os.environ['PYTHONNOUSERSITE'] = "yes"

    # When invoked by the system, the shell will be dash.
    workspace = "SHELL" in os.environ
    os.environ['SHELL'] = "/bin/bash"

    p_path, p_url = get_proxy_addr()
    cmdlist = [
              "jupyter", "notebook", "--no-browser", "--port=8000",
              "--ip=0.0.0.0",
              '--NotebookApp.base_url="%s"' % p_path
              ]

    token = get_cookie()[1]
    if token:
        cmdlist.append('--NotebookApp.token=%s' % token)

    if args.debug:
        cmdlist.append('--Application.log_level="DEBUG"')
    cmdlist += arglist

    urlstring = "%s?token=%s" % (p_url, token)
    print("=" * len(urlstring))
    print("URL is\n%s" % urlstring)
    print("=" * len(urlstring))
    sys.stdout.flush()  # make sure this gets displayed now.

    # If invoked by system, display logo for screenshots
    if not workspace:
        try:
            Popen(['xview', '-fit', '-fullscreen',
                  '/apps/jupyterexamples/current/images/jupyter.gif'])
        except:
            pass

    # having DISPLAY set breaks notebook terminals
    os.environ['DISPLAY'] = ""
    if args.debug:
        print("cmdlist:", cmdlist)
    notebookProcess = Popen(cmdlist)
#   print("%s notebookProcess = %d" % (time.ctime(),notebookProcess.pid))
#   sys.stdout.flush()  # make sure this gets displayed now.
    returnCode = notebookProcess.wait()
    notebookProcess = None
    print("\n\n%s Notebook server terminated. returnCode = %d\n" % (time.ctime(),returnCode))
    sys.stdout.flush()  # make sure this gets displayed now.


def find_header(file):
    # looks for HEADER in custom.js file. Returns version
    found = 0
    try:
        f = open(file, 'r')
        for line in f.read().split('\n'):
            if line.startswith('/** HEADER '):
                found = int(line.split()[2])
                break
        f.close()
    except:
        pass
    return found


header_template = """
/** HEADER 8 **/
require(["jquery", "base/js/namespace", "base/js/dialog"],
    function($, Jupyter, dialog) {
        var logo = document.getElementById("ipython_notebook").children[0];
        var nb_url = logo.href;
        var host = window.location.hostname.replace('proxy.', '');
        var uid = ${uid};
        try {
            var base = '/' + window.location.pathname.split('/').slice(1,5).join('/') + '/';
        } catch(err) {
            try {
                var base = Jupyter.notebook.base_url;
            } catch (err) {
                var base = IPython.notebook.base_url;
            }
        }
        var end = base.indexOf('/',7);
        var session = base.slice(7,end);

        try {
            const pattern = new RegExp('/tree$')
            var isTree = pattern.test(window.location.pathname)
            if (isTree) {
                var notificationMessagePath = "https://"+window.location.hostname+window.location.pathname.replace(pattern,"/tree/data/sessions/"+session+"/notificationMessage")
            };
        } catch(err) {
            var isTree = false
        }

        if (isTree) {
// based on: https://developer.mozilla.org/en-US/docs/Web/API/XMLHttpRequest/Synchronous_and_Asynchronous_Requests
            try {
                var xhr = new XMLHttpRequest();
                xhr.responseType = "text";
                xhr.timeout = 30000;
                xhr.ontimeout = function () {
                    console.error("The request for notificationMessage timed out.");
                };
                xhr.onloadend = function(pe) {
                    if (xhr.readyState === 4) {
                        if (xhr.status === 200) {
// based on: https://www.stefaanlippens.net/jupyter-notebook-dialog.html
                            dialog.modal({title: 'Notification',
                                          body: $(xhr.responseText),
                                          buttons: {'OK': {}}
                                         });
                        } else if (xhr.status === 0 && xhr.responseText != "") {
// This happens on older Jupyter Notebook tools
                            dialog.modal({title: 'Notification',
                                          body: $(xhr.responseText),
                                          buttons: {'OK': {}}
                                         });
                        } else if (xhr.statusText != "") {
                            console.error(xhr.statusText);
                        }
                    }
                };
                xhr.open('GET',notificationMessagePath,true);
                xhr.send(null);
            } catch(err) {
                console.error(err);
            }
        };

        var stat = 'https://' + window.location.host + base + 'static/';
        logo.href = 'https://' + host + '/members/' + uid + '/dashboard';
        logo.target = "_blank";
        logo.title = "${shost} Dashboard";
        var hostlogo = stat + '${lhost}.png';
        var jlogo = stat + 'jupyter.png';
        logo.innerHTML=`<img src="${hostlogo}" style="margin:0px 20px 0px 0px"><img src="${jlogo}">`;

        var term = document.getElementById("login_widget");
        var url = `https://${host}/tools/anonymous/stop?sess=${session}`;

        term.innerHTML = `<button class="btn btn-sm navbar-btn" title="Report a problem" onclick="window.open('https://${lhost}.org/feedback/report_problems')">Submit a ticket</button>&nbsp;&nbsp;<button class="btn btn-sm navbar-btn" title="Terminate this notebook or tool and any others in the session" onclick="window.location.href='${url}'">Terminate Session</button>`;

        try {
            var quit = document.getElementById("shutdown");
            quit.parentNode.removeChild(quit);
        } 
        catch(err) {
        }

        if (nb_url.indexOf('token') === -1) { return };
        url = '<li id=\"share\" title=\"Share Session\"><a href=\"#\" id=\"share\">Share Session</a></li>';
        $("#open_notebook").after(url);
        $('#share').off("click"); 
        $('#share').click(function () {
            txt = "The shareable URL below enables others to access your entire session. \
            Anyone with it can potentially read all your files on the hub. Be careful who you give it to.\
            Terminating this session removes their access. ";
            dialog.modal({
                title: 'Share Session',
                body: $('<div/>')
                    .text(txt).css('font-size', '24px')
                    .append($('<div/>')
                        .addClass('alert alert-info')
                        .text(nb_url).css('font-size', '12px')),
                buttons: {
                    OK: { 'class': 'btn-primary' }
                }
            });
        });
    }
);
"""

css_contents = """/****** AUTOMATICALLY GENERATED BY NotebookSettings **************/
/* Warning: Anything you add here will get overwritten! */
/** HEADER 8 **/

/* remove clusters tab in tree view */
.clusters_tab_link {
    visibility: hidden;
}
/* remove CellToobar Reminder */
[title="show new celltoolbar selector location"] {
    display: none;
}

/* FULL WIDTH */
.notebook_app .container {
    width:99%; !important
}

/* Prevent the edit cell highlight box from getting clipped;
 * important so that it also works when cell is in edit mode*/
div.cell.selected {
    border-left-width: 1px !important;
}
"""

def write_custom():
    cdir = os.path.expanduser("~/.jupyter/custom")
    if not os.path.exists(cdir):
        os.system('mkdir -p ' + cdir)

    cjs = os.path.join(cdir, 'custom.js')
    css = os.path.join(cdir, 'custom.css')

    # write js and css if necessary
    ver = find_header(cjs)
    if ver == 8:
        return

    with open(css, 'w') as f:
        f.write(css_contents)

    t = Template(header_template)
# hub_session_container
    host = os.uname()[1]
    host = host.split('_')[0]
    uid = os.getuid()
    with open(cjs, 'w') as f:
        f.write(t.safe_substitute(uid=uid, shost=host, lhost=host.lower()))


def cleanup():
    # misc cleanup and fixes

    # this can break nglview
    ndir = os.path.expanduser("~/.local/share/jupyter/nbextensions/nglview-js-widgets")
    if os.path.isdir(ndir):
        call(['rm', '-rf', ndir])


def main():
    if os.getuid() == 0:
        print("Do not run this as root.", file=sys.stderr)
        sys.exit(1)

    parser = parse_cmd_line()
    args = parser.parse_args()
    if args.help:
        parser.print_usage()
        sys.exit(0)

    write_custom()
    cleanup()
    run_notebook(args)


if __name__ == "__main__":