# Put in your SUNet ID.
# If your local username IS your SUNet ID, you can just enter "$USER" instead.
SUNETID=$USER

# Put in the hostname (JUST the hostname) of your group's bastion host
BASTION=luckdragon
#BASTION=dr

# If you want to change the name of the principal-switching commands,
# or you want to change the name of the SSH-related commands,
# look at the bottom of the script.

#
# CODE BEGINS HERE!
#

# Update the prompt with the current principal
get_principal_ps1() {
  klist -s
  if [ $? -eq 1 ]; then
    printf '(nc)'
    return 0
  fi

  krb5_username=$(klist | grep 'Default principal' | sed -e 's/Default principal: \(.*\)\@.*/\1/')

  if [ $krb5_username = $USER ]; then
    printf '(p)'
  elif [ $krb5_username = "$USER/root" ]; then
    printf '(R)'
  elif [ $krb5_username = "$USER/admin" ]; then
    printf '(-A-)'
  elif [ $krb5_username = "$USER/sunet" ]; then
    printf '(-S-)'
  else
    echo -n $krb5_username
  fi
  return 0
}
export PS1="\h:\w \u\$(get_principal_ps1)\\$ \[$(tput sgr0)\]"

# Support multiple Kerberos principals
# Takes one parameter, which is the type of principal to switch to
# (either empty string, 'root', 'admin', or 'sunet').
pag_switch() {
	# Check the PAG type we're switching to
	type=$1
	if [ "$type" = '' ]; then
		TARGET_PRINCIPAL="${SUNETID}@stanford.edu"
		FORWARDABLE=1
		type='regular'
	elif [ "$type" = 'root' ]; then
		TARGET_PRINCIPAL="${SUNETID}/root@stanford.edu"
		FORWARDABLE=0
	elif [ "$type" = 'admin' ]; then
		TARGET_PRINCIPAL="${SUNETID}/admin@stanford.edu"
		FORWARDABLE=0
	elif [ "$type" = 'sunet' ]; then
		TARGET_PRINCIPAL="${SUNETID}/sunet@stanford.edu"
		FORWARDABLE=0
	else
		echo "Invalid pag type $type"
		return
	fi

	# Some credentials can't be forwarded
	if [ $FORWARDABLE = 1 ]; then
		KINIT_OPTIONS='-f'
	else
		KINIT_OPTIONS='-F'
	fi
	KINIT_COMMAND="kinit $KINIT_OPTIONS $TARGET_PRINCIPAL"

	# Try switching
	echo "Switching to $TARGET_PRINCIPAL"
	need_kinit=0
	kswitch -p $TARGET_PRINCIPAL 2> /dev/null
	if [ $? -eq 1 ]; then
		need_kinit=1
	else
		klist -s
		if [ $? -eq 1 ]; then
			need_kinit=1
		fi
	fi

	# If we need to kinit, do it now
	if [ $need_kinit -eq 1 ]; then
		$KINIT_COMMAND
	fi
}

# SSH functions for connecting to Stanford systems.
ssh_stanford() {
	ssh -l $SUNETID $1.stanford.edu
}
ssh_stanford_root() {
	ssh -l root $1.stanford.edu
}
ssh_stanford_rootr() {
	ssh -l root.$SUNETID $1.stanford.edu
}
ssh_stanford_roota() {
	ssh -l admin.$SUNETID $1.stanford.edu
}

#
# CODE ENDS
# 

# Principal-switching aliases
# You can change the alias names to whatever you like, just so long as
# everything to the right of the = is unchanged.
alias pag='pag_switch'
alias rootpag='pag_switch root'
alias adminpag='pag_switch admin'
alias sunetpag='pag_switch sunet'

# SSH aliases
alias sshs=ssh_stanford
alias sshr=ssh_stanford_root
alias sshrr=ssh_stanford_rootr
alias sshra=ssh_stanford_roota
alias bastion='pag; exec ssh -p 44 -l $SUNETID $BASTION.stanford.edu bin/luckdragon.sh'