#!/bin/bash
#
# Obtain and store condor credentials for each given oauth_service.
# A single storer services credentials of multiple kinds, detecting the kind
# of each service from the configured provider-name lists (Vault via htgettoken,
# Pelican via the pelican client, ...).
# Can be run automatically from condor_submit when configured as
#   SEC_CREDENTIAL_STORER.
#
# Formerly named condor_vault_storer; that name is still installed as a symlink
# for backward compatibility.

ME="${0##*/}"

ARGS=("$@")

# When condor_submit launches this the cursor is not at the beginning
#   of a new line, so always send a newline before anything else to
#   stdout or stderr.
NEWLINENEEDED=true
newlineifneeded()
{
    if $NEWLINENEEDED; then
        echo
        NEWLINENEEDED=false
    fi
}

usage()
{
    newlineifneeded
    echo "Usage: $ME [-vd] oauth_service ..."
    echo "  -v shows more progress than normal" 
    echo "  -d shows a lot of debug information" 
    echo "  Each oauth_service is an issuer optionally followed by underscore and role"
    echo "  Options may be added to each oauth_service:"
    echo "    &handle=<name>       A name to be added to the credential"
    echo "    &scope=<scopes>      A comma-separated list of scopes to request"
    echo "    &audience=<audience> A token audience to request"
    echo "    &options=<args>      Extra arguments to htgettoken"
    exit 1
} >&2

fatal()
{
    newlineifneeded
    echo "$ME: $@"
    if [ "$VERBOSE" = "" ]; then
        echo "  More details might be available by running"
        echo "    $ME -v $(printf '"%s" ' "${ARGS[@]}")"
    fi
    exit 1
} >&2

verbose()
{
    if [ -n "$VERBOSE" ]; then
        newlineifneeded
        echo "$@"
    fi >&2
}

CONDOROPTS="`condor_config_val SEC_CREDENTIAL_VAULT_STORER_OPTS 2>/dev/null`"
VERBOSE=
parseopts()
{
    unset OPTIND
    while getopts "dv" opt; do
        case ${opt} in
            d) VERBOSE="-d";;
            v) VERBOSE="-v";;
            \?) usage;;
        esac
    done
}
parseopts $CONDOROPTS
parseopts "$@"
shift $((OPTIND -1))

if [ "$#" = 0 ]; then
    usage
fi

STOREOUT=/dev/null
if [ "" != "$VERBOSE" ]; then
    STOREOUT=/dev/stdout
fi

CONDOROPTS="`condor_config_val SEC_CREDENTIAL_GETTOKEN_OPTS 2>/dev/null`"
# The htgettoken options are only required for Vault-managed services.  A single
# storer may serve Vault and Pelican services side by side (and ignore services
# owned by the local/oauth2/client credmons), so defer this check until we
# actually need to fetch a Vault token below.
GETTOKENOPTS_SET=true
if [ -z "$CONDOROPTS" ] && [ -z "$HTGETTOKENOPTS" ]; then
    GETTOKENOPTS_SET=false
fi

# Provider lists used to decide which credmon owns each requested service.
PELICANPROVIDERS="`condor_config_val PELICAN_CREDMON_PROVIDER_NAMES 2>/dev/null`"
LOCALPROVIDERS="`condor_config_val LOCAL_CREDMON_PROVIDER_NAMES 2>/dev/null`"
if [ -z "$LOCALPROVIDERS" ]; then
    LOCALPROVIDERS="`condor_config_val LOCAL_CREDMON_PROVIDER_NAME 2>/dev/null`"
fi
OAUTH2PROVIDERS="`condor_config_val OAUTH2_CREDMON_PROVIDER_NAMES 2>/dev/null`"
CLIENTPROVIDERS="`condor_config_val CLIENT_CREDMON_PROVIDER_NAMES 2>/dev/null`"

# The Pelican client used to mint subject tokens for Pelican services; override
# with the PELICAN_BIN environment variable.  It must not require a TTY since it
# runs as a non-interactive subprocess of condor_submit.
PELICAN="${PELICAN_BIN:-pelican}"
export PELICAN_SKIP_TERMINAL_CHECK=1

# make sure we can find condor_store_cred
PATH=/usr/sbin:$PATH

# CONDOR_VAULT_STORER_USER and CONDOR_VAULT_STORER_ID can be set by a
# service that needs to impersonate different users, assuming the service
# has tokens for the users.
STOREUSER=""
if [ -n "$CONDOR_VAULT_STORER_USER" ]; then
    STOREUSER=" -u $CONDOR_VAULT_STORER_USER"
fi

# Keep the standard duration vault token in $VTOKEN-$SERVICE to make
#  sure that credmon has a long-duration one, but copy it to $VTOKEN if a
#  new one is generated.
ID="${CONDOR_VAULT_STORER_ID:-u`id -u`}"
VTOKEN="/tmp/vt_$ID"
BTOKEN=""
if [ -z "$BEARER_TOKEN_FILE" ]; then
    # Also store the bearer token with a -$SERVICE suffix
    BTOKEN="${XDG_RUNTIME_DIR:-/tmp}/bt_$ID"
fi

NL="
"

# These two functions depend on some variables that are set in the loop below,
# but they do not need to be defined multiple times so they are outside of the
# loop.
querycred() {
    STORENAME="$1"
    verbose "Checking if $@ credentials exist"
    STOREMSG="`condor_store_cred query-oauth$STOREUSER -s "$@" >$STOREOUT 2>&1`"
    case $? in
        0) ;;
        1)  if [ -n "$STOREMSG" ]; then
                verbose "$STOREMSG"
            fi
            return 1
            ;;
        2) fatal "Credentials exist that do not match the request.
They can be removed by
  condor_store_cred delete-oauth$STOREUSER -s $STORENAME
but make sure no other job is using them."
            ;;
        *) fatal "${STOREMSG}${NL}Querying condor credentials failed";;
    esac
}

storecred() {
    STORENAME="$1"

    if $SHOWSTORING; then
        echo "Storing condor credentials for $STORENAME" >&2
    fi
    (
    echo "{"
    if $STOREVAULTTOKEN; then
        echo "  \"vault_token\": \"$CRED\","
    fi
    echo "  \"vault_url\": \"$VAULTURL\""
    echo "}"
    )|condor_store_cred add-oauth$STOREUSER -s "$@" -i - >$STOREOUT
    if [ $? != 0 ]; then
        fatal "Failed to store condor credentials for $STORENAME"
    fi
}

# ---------------------------------------------------------------------------
# Pelican service support
#
# A Pelican service is owned by the PelicanCredmon (listed in
# PELICAN_CREDMON_PROVIDER_NAMES).  Instead of htgettoken/Vault, we run the
# Pelican client's device-code flow to obtain a subject access token and store
# it as the service's .top credential; the PelicanCredmon then performs the
# RFC 8693 token exchange and ongoing refreshes.
# ---------------------------------------------------------------------------

# service_in_list <name> <list>: succeed if <name> appears in the whitespace or
# comma separated <list>.
service_in_list() {
    local name="$1" list="$2" item
    local IFS=', '
    for item in $list; do
        if [ "$item" = "$name" ]; then
            return 0
        fi
    done
    return 1
}

store_pelican_service() {
    local REQUEST="$1" SERVICE="$2"
    local -a PPARTS
    IFS="&" read -r -a PPARTS <<< "$REQUEST"
    local HANDLE="" REQ_SCOPE="" REQ_AUDIENCE="" PART VAL
    for PART in "${PPARTS[@]:1}"; do
        VAL="${PART#*=}"
        case "$PART" in
            handle=*)         HANDLE="$VAL" ;;
            scope=*|scopes=*) REQ_SCOPE="$VAL" ;;
            audience=*)       REQ_AUDIENCE="$VAL" ;;
        esac
    done

    local STORENAME="$SERVICE"
    if [ -n "$HANDLE" ]; then
        STORENAME="${SERVICE}_${HANDLE}"
    fi

    # If the credd already holds a credential for this service, the credmon
    # keeps it refreshed; nothing to do.
    if condor_store_cred query-oauth$STOREUSER -s "$STORENAME" >/dev/null 2>&1; then
        verbose "Credentials for $STORENAME already stored; nothing to do"
        return
    fi

    if ! command -v "$PELICAN" >/dev/null 2>&1; then
        fatal "Pelican client '$PELICAN' not found (set PELICAN_BIN)"
    fi

    local PELICAN_URL PREFIX PERMS
    PELICAN_URL="`condor_config_val ${SERVICE}_PELICAN_URL 2>/dev/null`"
    PREFIX="`condor_config_val ${SERVICE}_PELICAN_PREFIX 2>/dev/null`"
    PERMS="`condor_config_val ${SERVICE}_PELICAN_PERMISSIONS 2>/dev/null`"
    if [ -z "$PERMS" ]; then
        PERMS="read"
    fi
    if [ -z "$PELICAN_URL" ]; then
        fatal "Service $SERVICE: ${SERVICE}_PELICAN_URL is not configured"
    fi

    # Build the pelican:// resource URL from the federation URL + prefix.
    local RESOURCE="$PELICAN_URL"
    case "$PELICAN_URL" in
        pelican://*|osdf://*) : ;;
        https://*) RESOURCE="pelican://${PELICAN_URL#https://}" ;;
        http://*)  RESOURCE="pelican://${PELICAN_URL#http://}" ;;
    esac
    if [ -n "$PREFIX" ]; then
        RESOURCE="${RESOURCE%/}/${PREFIX#/}"
    fi

    # The permissions are a list (whitespace and/or comma separated); a token
    # may carry several capabilities at once, since they do not imply one
    # another (e.g. modify does not imply read).  Map each to a Pelican client
    # capability flag.
    local -a FETCH_FLAGS=()
    local PERM
    for PERM in ${PERMS//,/ }; do
        case "$PERM" in
            read)          FETCH_FLAGS+=("--read") ;;
            write|create)  FETCH_FLAGS+=("--write") ;;
            modify|delete) FETCH_FLAGS+=("--modify") ;;
            "") ;;
            *) fatal "Service $SERVICE: unsupported permission '$PERM' (use read, write, and/or modify)" ;;
        esac
    done
    if [ ${#FETCH_FLAGS[@]} -eq 0 ]; then
        fatal "Service $SERVICE: ${SERVICE}_PELICAN_PERMISSIONS lists no usable permission"
    fi

    newlineifneeded
    echo "$ME: Obtaining a Pelican token for service '$SERVICE' ($RESOURCE, $PERMS)" >&2

    # Run the device-code flow.  The token is printed to stdout; the "navigate
    # to this URL" prompt goes to stderr (relayed to the user by condor_submit).
    local SUBJECT_TOKEN
    SUBJECT_TOKEN="`"$PELICAN" token fetch "${FETCH_FLAGS[@]}" "$RESOURCE"`"
    if [ $? != 0 ] || [ -z "$SUBJECT_TOKEN" ]; then
        fatal "Failed to obtain a Pelican token for service '$SERVICE'"
    fi
    # Keep only the JWT (last non-empty line) in case extra text is printed.
    SUBJECT_TOKEN="`printf '%s\n' "$SUBJECT_TOKEN" | awk 'NF{line=$0} END{print line}'`"

    local -a PSTOREOPTS=()
    if [ -n "$REQ_SCOPE" ]; then
        PSTOREOPTS+=("-S" "$REQ_SCOPE")
    fi
    if [ -n "$REQ_AUDIENCE" ]; then
        PSTOREOPTS+=("-A" "$REQ_AUDIENCE")
    fi

    if $SHOWSTORING; then
        echo "Storing condor credentials for $STORENAME" >&2
    fi

    # Hand the subject token to the credd as the .top credential.  The
    # PelicanCredmon discovers the issuer's token endpoint from this token, so we
    # record only the token itself.  No "vault_url" is included so the credd does
    # not treat it as a Vault token.
    {
        printf '{ "access_token": "%s" }\n' "$SUBJECT_TOKEN"
    } | condor_store_cred add-oauth$STOREUSER -s "$STORENAME" "${PSTOREOPTS[@]}" -i - >$STOREOUT
    if [ $? != 0 ]; then
        fatal "Failed to store Pelican credentials for '$STORENAME'"
    fi
    verbose "Stored subject token for $STORENAME"
}

for REQUEST; do
    SHOWSTORING=false
    if [ -n "$VERBOSE" ]; then
        SHOWSTORING=true
    fi
    IFS="&" read -r -a PARTS <<< "$REQUEST"
    SERVICE="${PARTS[0]}"
    ISSUER="${SERVICE%%_*}"

    # Route each requested service to the credmon that owns it, so a single
    # storer can serve a mix of Vault and Pelican providers on one access point.
    #
    # condor_submit tells us which kind of credential these services are via
    # CONDOR_CREDENTIAL_STORER_MODE (it groups a submit's services by kind and
    # runs us once per kind).  Trust that when set; otherwise -- e.g. when
    # invoked by an older condor_submit, or by hand -- fall back to detecting the
    # owner from the configured provider-name lists.
    SERVICEMODE="$CONDOR_CREDENTIAL_STORER_MODE"
    if [ -z "$SERVICEMODE" ]; then
        if service_in_list "$SERVICE" "$PELICANPROVIDERS" || \
           service_in_list "$ISSUER" "$PELICANPROVIDERS"; then
            SERVICEMODE=pelican
        elif service_in_list "$SERVICE" "$LOCALPROVIDERS" || \
             service_in_list "$ISSUER" "$LOCALPROVIDERS" || \
             service_in_list "$SERVICE" "$OAUTH2PROVIDERS" || \
             service_in_list "$SERVICE" "$CLIENTPROVIDERS"; then
            SERVICEMODE=other
        else
            SERVICEMODE=vault
        fi
    fi

    case "$SERVICEMODE" in
    pelican)
        store_pelican_service "$REQUEST" "$SERVICE"
        continue
        ;;
    other)
        # Managed by the local/oauth2/client credmons; never delivered through
        # this storer.  Skip defensively.
        verbose "Service $SERVICE is managed by another credmon; skipping"
        continue
        ;;
    esac

    # Otherwise this is a Vault-managed service, which needs htgettoken options.
    if ! $GETTOKENOPTS_SET; then
        fatal 'Neither SEC_CREDENTIAL_GETTOKEN_OPTS condor value nor $HTGETTOKENOPTS environment set'
    fi

    # using arrays for options works better for quoting
    read -r -a OPTS <<< "$CONDOROPTS"
    OPTS+=("-i" "$ISSUER")
    if [ "$SERVICE" != "$ISSUER" ]; then
        ROLE="${SERVICE#*_}"
        # With ISSUER as everything before the first underscore and ROLE as
        # everything after it, keep enforcing the historical rule that a Vault
        # service name has at most one underscore: a ROLE containing another
        # underscore means the name was malformed.
        if [ "$ROLE" != "${ROLE#*_}" ]; then
            fatal "Only one underscore allowed in use_oauth_services name"
        fi
        OPTS+=("-r" "$ROLE")
    fi
    STOREOPTS=()
    if [ -n "$SEC_CREDENTIAL_STORECRED_OPTS" ]; then
        for OPT in $SEC_CREDENTIAL_STORECRED_OPTS; do
            STOREOPTS+=($OPT)
        done
    fi
    HANDLE=""
    HANDLESERVICE=""
    HANDLESTOREOPTS=("${STOREOPTS[@]}")
    HANDLEOPTS=("${OPTS[@]}")
    REDUCED=false
    for PART in "${PARTS[@]:1}"; do
        VAL="${PART#*=}"
        case "$PART" in
            handle=*)
                HANDLE="$VAL"
                HANDLESERVICE="${SERVICE}_$HANDLE"
                ;;
            scopes=*)
                HANDLESTOREOPTS+=("-S" "$VAL")
                HANDLEOPTS+=("--scopes=$VAL")
                REDUCED=true
                ;;
            audience=*)
                HANDLESTOREOPTS+=("-A" "$VAL")
                HANDLEOPTS+=("--audience=$VAL")
                REDUCED=true
                ;;
            options=*)
                OPTS+=($VAL)
                HANDLEOPTS+=($VAL)
                ;;
        esac
    done

    if $REDUCED && [ -z "$HANDLE" ]; then
        fatal "reduced scopes or audience require a handle"
    fi

    if [ -n "$BTOKEN" ]; then
        OPTS=("-o" "$BTOKEN-$SERVICE" "${OPTS[@]}")
    fi
    OPTS=("--vaulttokenttl=28d" "${OPTS[@]}" "--vaulttokeninfile=$VTOKEN-$SERVICE" "--vaulttokenfile=/dev/stdout" "--showbearerurl")

    # This next section may need to be done twice if it turns out that
    # there is a valid $VTOKEN-$SERVICE vault token but there is no 
    # corresponding credential stored in credd.  The query of credd 
    # may require an access token, however, so the query can't be done until
    # after the first time htgettoken runs.
    TRIEDAGAIN=false
    while true; do 
        verbose "Attempting to get tokens for $SERVICE"
        # First attempt to get tokens quietly without oidc.
        # If a valid vaulttokeninfile exists, it will not generate a new vault
        # token; the 28 day vault token will only be created (in stdout) if the
        # vaulttokeninfile is not there or expired.
        CRED="`htgettoken "${OPTS[@]}" --nooidc ${VERBOSE:--q}`"
        if [ $? != 0 ]; then
            # OIDC authentication probably needed, so remove -q to tell the user
            #  what is happening
            SHOWSTORING=true
            newlineifneeded
            echo "Authentication needed for $SERVICE" >&2
            CRED="`htgettoken "${OPTS[@]}" $VERBOSE`"
            if [ $? != 0 ]; then
                fatal "htgettoken failed"
            fi
        fi

        if [ -n "$BTOKEN" ] && [ -f $BTOKEN-$SERVICE ]; then
            verbose "Copying bearer token to $BTOKEN"
            # Copy bearer token to $BTOKEN atomically
            TMPFILE="`mktemp $BTOKEN.XXXXXXXXXX`"
            cat $BTOKEN-$SERVICE >$TMPFILE
            mv $TMPFILE $BTOKEN
        fi

        STOREVAULTTOKEN=false
        VAULTURL=""
        CREDLINES="`echo "$CRED"|wc -l`"
        case $CREDLINES in
            1)  # No new vault token was generated
                # There is only a vault url, which we will not need

                if $TRIEDAGAIN; then
                    fatal "No new vault token after second try; logic error"
                fi
                if ! querycred $SERVICE "${STOREOPTS[@]}"; then
                    if [ -f $VTOKEN-$SERVICE ] && [ ! -w $VTOKEN-$SERVICE ]; then
                        fatal "No $SERVICE credentials stored and $VTOKEN-$SERVICE is readonly"
                    fi
                    verbose "Removing $VTOKEN-$SERVICE because there are no $SERVICE credentials stored"
                    rm -f "$VTOKEN-$SERVICE"
                    TRIEDAGAIN=true
                    continue
                fi

                # Do not need to create new vault token
                CRED=""
                ;;
            2)  # First line is new long-lived vault token, second line is vault url
                VAULTURL="`echo "$CRED"|(read X; read L; echo "$L")`"
                CRED="`echo "$CRED"|(read L; echo "$L")`"
                STOREVAULTTOKEN=true
                ;;
            *)  fatal "Unexpected number of stdout lines from htgettoken: $CREDLINES";;
        esac
        break
    done

    if [ -z "$CRED" ] && [ -z "$HANDLE" ]; then
        continue
    fi

    # $CRED now either contains a new, long duration vault token which
    # needs to be exchanged for a shorter duration one, or $CRED is empty
    # and we just need to invoke htgettoken to calculate the vault URL that
    # credmon will need to use with the handle.  The above invocation does
    # not include reduced scopes and audiences associated with the handle,
    # so we couldn't have htgettoken calculate that URL there.
    # Normally do this exchange quietly.
    TOKENOUT=
    if [ -n "$HANDLE" ] && [ -n "$BTOKEN" ]; then
        # We don't necessarily need to save this token, since really the
        # weakened token is only needed in the job, but since we're generating
        # it anyway we might as well save it.  Consider stopping to generate
        # it when --nobearertoken works with --showbearerurl.
        TOKENOUT="$BTOKEN-$HANDLESERVICE"
    else
        # TODO: When this case happens and after htgettoken 2.0 or greater
        # is installed everywhere, use --nobearertoken instead of using a
        # temporary token name.  In older versions --nobearertoken didn't work
        # properly in combination with --showbearerurl.
        TOKENOUT="`mktemp ${XDG_RUNTIME_DIR:-/tmp}/bt_tmp.XXXXXXXXXX`"
        trap "rm -f $TOKENOUT" 0
    fi

    if [ -n "$HANDLE" ]; then
        verbose "Weakening token for ${HANDLESERVICE}"
    fi

    HANDLEURL=
    EXCHANGEOPTS="--nooidc --nokerberos --nossh -o $TOKENOUT --vaulttokenfile=$VTOKEN-$SERVICE --showbearerurl ${VERBOSE:--q}"
    if [ -n "$CRED" ]; then
        HANDLEURL="`echo "$CRED"|htgettoken "${HANDLEOPTS[@]}" --vaulttokeninfile=/dev/stdin $EXCHANGEOPTS`"
        if [ $? != 0 ]; then
            fatal "Failed to exchange vault token"
        fi
    else
        # Note that this case only happens when $HANDLE is set, although
        # it's possible for $HANDLE to be set with the above case as well.
        HANDLEURL="`htgettoken "${HANDLEOPTS[@]}" $EXCHANGEOPTS`"
        if [ $? != 0 ]; then
            fatal "Failed to obtain weakened token"
        fi
    fi

    verbose "Copying $VTOKEN-$SERVICE to $VTOKEN"
    TMPFILE="`mktemp $VTOKEN.XXXXXXXXXX`"
    cat $VTOKEN-$SERVICE >$TMPFILE
    mv $TMPFILE $VTOKEN

    if $STOREVAULTTOKEN; then
        storecred $SERVICE "${STOREOPTS[@]}"
    fi

    if [ -n "$HANDLE" ]; then
        if  $STOREVAULTTOKEN || ! querycred $HANDLESERVICE "${HANDLESTOREOPTS[@]}"; then
            # There's not really a need to store this because of
            #  $STOREVAULTTOKEN but that's rare enough that we might
            #  as well, just in case the URL changes or something.
            VAULTURL="$HANDLEURL"
            storecred $HANDLESERVICE "${HANDLESTOREOPTS[@]}"
        fi
    fi
done
