Interaction with SWAN APIs from a script

import requests
import subprocess
from auth_get_sso_cookie import cern_sso
import pprint
import logging
import re
import time

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# MAIN PARAMETERS
USER = '<NICE username>'
PASSWORD = '<NICE password>'
AUTH_SERVER = "auth.cern.ch"
NOCERTVERIFY = False

# GUESS THE SWAN MACHINE
r = requests.get('https://swan.cern.ch')
# Find the name of the machine from the html response of the main page
groups = re.findall("window\.location\.replace\(\\\'(https://swan\d+\.cern\.ch).+$", r.text, re.MULTILINE)
if not len(groups):
    raise Exception("Regex didn't find the machine to connect to")

MACHINE = groups[0]
OAUTH_LOGIN = f'{MACHINE}/hub/oauth_login'
print(f'Machine: {MACHINE}')

# CALL EXTERNAL PROGRAMS (KERBEROS)
# Get keytab file for user and ticket from kerberos
keytab_command = subprocess.run(f'cern-get-keytab --user --login {USER} --password {PASSWORD} --keytab {USER}.keytab'.split())
kdestroy_command = subprocess.run('kdestroy')
kinit_command = subprocess.run(f'kinit -f -r 5d -kt {USER}.keytab {USER}'.split())
klist_command = subprocess.run(['klist'])
print(klist_command.stdout, flush=True)

# SESSIONS AND COOKIES ACQUISITION
# Create a session to store cookies
session, response = cern_sso.login_with_kerberos(OAUTH_LOGIN, not NOCERTVERIFY, AUTH_SERVER, silent=False)
if response.status_code == 302:
    redirect_uri = response.headers["Location"]
else:
    raise Exception('Redirect uri missing from "Location" header of response')
# Get a request to the redirect uri to get jupyterhub's cookies
session.get(redirect_uri, verify=not NOCERTVERIFY)
# BE SURE TO USE REFERER OR YOU'LL GET 403
session.headers['Referer'] = f'{MACHINE}/hub/token'
# Now you should have acquired the cookies to access the JH APIs
for key in ['jupyterhub-hub-login', 'jupyterhub-session-id']:
    print(key, session.cookies.get(key, 'None'))

# TEST THE APIS
# Generate a JH token 
r = session.post(f'{MACHINE}/hub/api/users/{USER}/tokens',
json={'notes': 'Token generated from script'})
print(f'Generation of token: {r.status_code} {r.json()}')

# Get User model in JupyterHub
r = session.get(f'{MACHINE}/hub/api/users/{USER}')
print(f'User model: {r.json()}')

# Stop a single-user server
r = session.delete(f'{MACHINE}/hub/api/users/{USER}/server')
if r.status_code == 204:
    print("The user's notebook server has stopped")
elif r.status_code == 202:
    print("The user's notebook server has not yet stopped as it is taking a while to stop")

# Start a single-user server
spawn_options = {'LCG-rel': 'LCG_97apython3',
'platform': 'x86_64-centos7-gcc8-opt',
'scriptenv': 'None',
'ncores': 2,
'memory': '8G',
'spark-cluster': 'none'}
print(f'Starting a single-user server with spawn options: {spawn_options}')
r = session.post(f'{MACHINE}/hub/api/users/{USER}/server', json=spawn_options)
print(f'Response from server: {r.content}')
if r.status_code == 201:
    print("The user's notebook server has started")
elif r.status_code == 202: 
    while r.status_code == 202:
        print("The user's notebook server has not yet started, but has been requested. Waiting for it to be ready")
        time.sleep(5)
        r = session.get(f'{MACHINE}/user/{USER}/api/status')
        json_response = r.json()
        print(r.status_code, json_response)
elif r.status_code == 500:
    print('Server returned 500. Something bad happened')
    

# Try to use jupyter server api.
# From https://petstore.swagger.io/?url=https://raw.githubusercontent.com/jupyter/notebook/master/notebook/services/api/api.yaml
# Ex. list content of directory "SWAN_projects"
t = session.get(f'{MACHINE}/user/{USER}/api/contents/SWAN_projects')
print('Folder contents using Jupyter APIs: ')
pprint.pprint(t.json())

# Stop the server to clean up
r = session.delete(f'{MACHINE}/hub/api/users/{USER}/server')
if r.status_code == 204:
    print("The user's notebook server has stopped")
elif r.status_code == 202:
    print("The user's notebook server has not yet stopped as it is taking a while to stop")