SLAC simple batch interface with kerberos authentication

client.py

import requests
from requests_kerberos import HTTPKerberosAuth
import os
from os.path import join  # Used to build url, assumes non-windows

# If this is a service, Start the process with 
# /usr/local/bin/pagsh k5start -t -f [/etc/path/to/keytab.keytab] [[email protected]] python

URL = os.environ("SLATCH_HOST") # i.e. "http://scalnx-v01.slac.stanford.edu:5000"
PRINCIPAL = os.environ("SLATCH_PRINCIPAL")
auth = HTTPKerberosAuth(principal=PRINCIPAL) # assumes service="http"

r = requests.get(url ,auth=auth).json()                          # GET /
r = requests.get(join(url, "jobs") ,auth=auth).json()            # GET /jobs
r = requests.get(join(url, "jobs", "1234") ,auth=auth).json()    # GET /jobs/1234

run.sh

#!/bin/bash
export SLATCH_HOSTNAME=example.slac.stanford.edu
# Need KRB5_KTNAME to point to SLAC-issued keytab with HTTP service for $SLATCH_HOSTNAME
export KRB5_KTNAME=/etc/krb5.keytab

FLASK_APP=slatch.py flask run --host=$SLATCH_HOSTNAME

slatch.py

import flask
from flask import Flask, make_response
from flask_kerberos import init_kerberos
from flask_kerberos import requires_authentication
from werkzeug.contrib.cache import SimpleCache

from functools import wraps
import getpass
import json
from socket import getfqdn

DEBUG=True

app = Flask(__name__)
init_kerberos(app, service="HTTP", getfqdn()) 
cache = SimpleCache()

STATUSES = "STATUSES"
USER = getpass.getuser()


@app.route('/')
@requires_authentication
def route_root(user):
    return "User '{}' is authenticated".format(user)


@app.route('/jobs', methods=['GET'])
@requires_authentication
def get_status(user):
    statuses = cached_statuses()
    return json.dumps(statuses.values())


@app.route('/jobs/<string:job_id>', methods=['GET'])
@requires_authentication
def get_statuses(user):
    statuses = cached_statuses()
    return json.dumps(statuses[job_id]) if job_id in statuses else make_response("Job Not Found", 404)


@app.route("/jobs", methods=['POST'])
@requires_authentication
def create_job(user):
    """
    Should return a new job_id
    """
    # build and submit job using bsub
    pass


@app.route('/jobs/<string:job_id>', methods=['PATCH'])
@requires_authentication
def modify_job(user):
    new_state = request.args.get("state")
    statuses = uncached_statuses()
    # alter_job not implemented
    return alter_job(job_id, state) if job_id in statuses else make_response("No job found", 404)


@app.route('/files', methods=['POST', 'PUT'])
@requires_authentication
def create_or_update_file(user):
    """
    If POST, throw error if file exists. If PUT, then always write file. 
    """
    pass


@app.route('/files', methods=['GET'])
@requires_authentication
def get_file(user):
    pass


def cached_statuses():
    statuses = cache.get(STATUSES)
    if not statuses:
        statuses = uncached_statuses()
        cache.set(STATUSES, statuses, timeout=60)
    return statuses


def uncached_statuses():
    args = STATUS_COMMAND.split()
    args.append(getpass.getuser())
    process = subprocess.Popen(args, stdout=subprocess.PIPE)
    print process.stdout.readline()  # Consume header
    stats = {}
    for line in process.stdout:
        line = line[:-1]         # omit newline
        if not line:             # omit blank lines
            continue
        if line[0] == " ":  # Post-process comment lines
            stats[-1]["comment"] = line
            continue
        # Translate '-' columns to None
        match = [i if i != '-' else None for i in line.split()]  
        stat = {}
        job_id = match[0]
        stat["job_id"] = job_id
        stat["user"] = match[1]
        stat["status"] = match[2]
        stat["queue"] = match[3]
        stat["host"] = match[5]
        stat["submitted"] = match[7]
        stat["started"] = match[13] if len(match) > 13 else None
        stat["ended"] = match[14] if len(match) > 14 else None
        stat["cpu_used"] = match[9]
        stat["memory_used"] = match[10]
        stat["swap_used"] = match[11]
        stats[job_id] = stat
        # Might be good to get working directory and logfile too
    rc = process.wait()
    return stats