#! /bin/sh
#
# Copyright (c) 2013 Twitter, Inc.
#
# Originally written by Jan Schaumann <jschauma@twitter.com> in April 2013.
#
# This script allows you to easily share secrets with other users in your
# organization by way of ssh pubkeys.

# Set these two variables to match your use of LDAP for ssh keys.
#LDAPFIELD="SSHPubkey"
#LDAPSEARCH="ldapsearch -LLLxh ldap.yourdomain.com -b dc=example,dc=com"

### You should not have to edit anything below this line.

set -e

###
### Globals
###

umask 077

VERSION="1.8.1"
CAN_DECRYPT_FROM="1.8 1.7 1.6 1.5 1.4"
CAN_ENCRYPT_FOR="1.8 1.7 1.6 1.5"

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

ACTION="encrypt"
URLS=""
EXPERT="no"
FILE="/dev/stdin"
GITHUB="https://github.com/<user>.keys"
KEYFILES=""
RECIPIENTS=""
RECIPIENT_GROUPS=""
UNAME="$(uname -s)"
BASE64="base64"

TDIR="$(mktemp -d "${TMPDIR:-/tmp}/jass.XXXX")"
LDAP_KEYFILE="${TDIR}/ldap"
KEYDIR="${TDIR}/pubkeys"
PEMDIR="${TDIR}/pems"
SESSION_KEY="${TDIR}/secret"
URLS_KEYFILE="${TDIR}/urls"

VERBOSITY=0

###
### Functions
###

# Read one ssh key per line from the given file and append it to the
# 'keyfile'.  This will strip any key options and comments, then add the
# given recipient name as a comment to the file.
appendKey() {
	local readonly keyfile="${TDIR}/keyfile"

	local readonly file="${1}"
	local recipient="${2}"

	if [ ! -s "${file}" ]; then
		return
	fi

	oIFS="${IFS}"
	IFS='
'

	# If 'file' and 'recipient' are identical, then we were given a
	# pathname instead of the keys of a user.  We then provide the
	# basename of the path as an identifying comment.
	if [ x"${file}" = x"${recipient}" ]; then
		recipient="${recipient##*/}"
	fi

	verbose "Appending keys for '${recipient}' to keyfile'..." 3

	# SSH keys are more complex than you think.  They may contain
	# leading options (including quoted strings containing spaces) and
	# trailing comments.  All we know -- and care about -- are the
	# actual key.  Since we only support RSA keys, we can ignore the
	# other options; if somebody creates a key with this pattern as a
	# comment, then they can go screw themselves.
	for line in $(sed -n -e 's/.*\(ssh-rsa AAAAB3NzaC1[^ ]*\).*/\1/p' "${file}"); do
		echo "${line} ${recipient}" >> "${keyfile}"
	done

	IFS="${oIFS}"
}


# Let us pretend!  Versions are always either M.m or M.m.t.  We claim that
# tiny version numbers would not introduce any incompatibility issues, so
# we strip that.  This allows us to have tiny version increments without
# having to explicitly list every single tiny version in our compatibility
# list.
checkCompatibility() {
	local greater v v_maj_min
	local pass=0
	local readonly this_major_min=$(echo "${VERSION}" | sed -e 's/\([0-9]*\.[0-9]*\).*/\1/')

	verbose "Checking version compatibility..." 2

	if [ ! -f "${KEYDIR}/version" ]; then
		verbose "Skipping version compatibility test for input without version data." 3
		return
	fi

	for v in $(awk -F: '/CAN_ENCRYPT_FOR/ { print $2 }' "${KEYDIR}/version"); do
		v_maj_min=$(echo ${v} | sed -e 's/\([0-9]*\.[0-9]*\).*/\1/')
		greater=$(echo "${this_major_min} >= ${v_maj_min}" | bc)
		if [ "${greater}" -gt 0 ]; then
			return
		fi
	done

	echo "Incompatible jass format.  Input was generated using:" >&2
	cat "${KEYDIR}/version" >&2
	echo >&2
	echo "This is jass version ${VERSION}."
	exit 1
	# NOTREACHED
}


cleanup() {
	rm -fr "${TDIR}"
}

convertKeyFiles() {
	local pemfile pubkey
	local success=0

	verbose "Converting ssh pubkeys into PKCS8 format..." 2

	for pubkey in "${KEYDIR}"/*; do
		pemfile=$(basename ${pubkey})
		ssh-keygen -P '' -f "${pubkey}" -e -m PKCS8 > "${PEMDIR}/${pemfile}" 2>/dev/null ||	\
			verbose "Unable to convert ssh key '$(basename ${pubkey})' to PKCS8 format."
		if [ -s "${PEMDIR}/${pemfile}" ]; then
			success=1
		fi
	done

	if [ ${success} -eq 0 ]; then
		echo "No valid PKCS8 file could be generated." >&2
		exit 1
	fi
}

# split the single file containing all pubkeys into per-pubkey files
createKeyFiles() {
	local readonly keyfile="${TDIR}/keyfile"
	local found fp recipient type

	found=0

	verbose "Creating per-pubkey files..." 3
	(
		cd "${KEYDIR}"
		grep -v "^#" "${keyfile}" | grep . | sort -u | split -a 6 -l 1 - .key-
	)

	for k in "${KEYDIR}"/.key-*; do
		# We added the recipient as the last field in getPubkeys.
		recipient="$(awk '{print $NF}' "${k}")"
		type=$(ssh-keygen -e -f ${k} | sed -n -e 's/Comment:.*bit \(.*\),.*/\1/p')
		if [ x"${type}" != x"RSA" ]; then
			verbose "Key type of type '${type}' not supported. Skipping." 4
			continue
		fi

		fp=$(getFingerprintFromKey ${k} encrypt)
		if [ -n "${fp}" ]; then
			found=1
			mv ${k} "${KEYDIR}"/${recipient}-${fp}
		else
			verbose "Skipping invalid keyfile '${k}'..." 3
		fi
	done

	if [ ${found} -eq 0 ]; then
		echo "No supported key found." >&2
		exit 1
	fi
}

decrypt() {
	local output="/dev/stdout"
	local outfile skey skeyfp
	local f
	local found=0

	local readonly msgdir="${TDIR}/message"
	local readonly localfp=$(getFingerprintFromKey ${KEYFILES} decrypt)

	if [ -z "${localfp}" ]; then
		exit 1
	fi

	verbose "Decrypting..."

	verbose "Decoding data..." 2
	verbose "Splitting input..." 3
	cat "${FILE}" | ( cd "${TDIR}" && awk '/^begin-base64 / { 		\
							f="data"i++; }		\
						{ print > f; }			\
						/^====/ { close(f); }' )

	if [ $? -gt 0 ]; then
		echo "Invalid input?" >&2
		exit 1
	fi

	verbose "Uudecoding ecnrypted message..." 3
	( mkdir -p "${msgdir}" && cd "${msgdir}" && uudecode ../data0 )
	outfile=$(basename "${msgdir}"/*)

	verbose "Uudecoding all session keys..." 3
	if [ ! -s "${TDIR}"/data1 ]; then
		echo "No encoded session keys found." >&2
		exit 1
	fi
	( cd "${KEYDIR}" && uudecode ../data[1-9]* )

	if [ x"${EXPERT}" = x"no" ]; then
		checkCompatibility
	fi

	verbose "Identifying matching session key..." 3

	for f in "${KEYDIR}"/*; do
		skey=$(basename "${f}")
		skeyfp=$(echo ${skey} | sed -e 's/.*-//')

		if [ x"${skeyfp}" = x"${localfp}" ]; then
			verbose "Using ${skey}..." 4
			found=1
			break
		fi
	done

	if [ ${found} -eq 0 ]; then
		echo "Data was not encrypted for the key in ${KEYFILES}." >&2
		exit 1
	fi

	verbose "Decrypting session key..." 2
	openssl rsautl -inkey ${KEYFILES} -decrypt -in "${KEYDIR}/${skey}" -out "${TDIR}/skey" || {
		echo >&2
		echo "Unable to decrypt session key." >&2
		echo "Are you sure the passphrase unlocks ${KEYFILES}?" >&2
		exit 1
	}

	verbose "Decrypting message using session key..." 2
	openssl enc -d -aes-256-cbc -pass "file:${TDIR}/skey"	\
		-in "${msgdir}/${outfile}" -out "${output}"
}

encrypt() {
	local outfile="message"
	local fp
	local pemfile

	verbose "Encrypting..."

	if [ x"${FILE}" != x"/dev/stdin" ]; then
		if [ ! -r "${FILE}" ]; then
			echo "Unable to read '${FILE}'." >&2
			exit 1
		fi
		outfile=$(basename "${FILE}")
	fi

	getPubkeys
	createKeyFiles
	convertKeyFiles
	generateSessionKey

	verbose "Encrypting data with session key..." 2
	openssl enc -aes-256-cbc -salt -in "${FILE}" -pass "file:${SESSION_KEY}" | \
		uuencode -m "${outfile}"

	verbose "Encrypting session key with pubkeys..." 2
	for pemfile in "${PEMDIR}"/*; do
		fp=$(basename "${pemfile}")
		verbose "Encrypting session key with pubkey ${fp}..." 3
		openssl rsautl -pubin -inkey "${pemfile}" -encrypt -pkcs -in "${SESSION_KEY}" |	\
			uuencode -m ${fp}
	done

	verbose "Encoding version information..." 2
	( echo "VERSION: ${VERSION}";
	  echo "CAN_ENCRYPT_FOR: ${CAN_ENCRYPT_FOR}";
	  echo "CAN_DECRYPT_FROM: ${CAN_DECRYPT_FROM}";
	) | uuencode -m version
}

expandGroup() {
	local readonly group="${1}"
	local all group_users

	verbose "Expanding group '${group}'..." 2

	# Supplementary groups are easy...
	verbose "Expanding group '${group}' by supplementary group membership..." 3
	group_users="$(grep "^${group}:" /etc/group | awk -F: '{print $NF}' | tr ',' ' ')"
	all="${group_users}"
	if [ -n "${all}" ]; then
		verbose "${all}" 4
	fi

	# ...but some users have primary groups...
	gid=$(grep "^${group}:" /etc/group | awk -F: '{print $3}')
	if [ -n "${gid}" ]; then
		verbose "Expanding group '${group}' by primary group membership..." 3
		group_users="$(awk -v gid="${gid}" -F: '{ if ($4 == gid) { print $1; }}' /etc/passwd)"
		if [ -n "${group_users}" ]; then
			all="${group_users} ${all}"
			verbose "${all}" 4
		fi
	fi

	# ...and of course there's LDAP.
	if [ -z "${all}" ]; then
		if ldapVarCheck; then
			verbose "Expanding group '${group}' from LDAP..." 3
			group_users="$(${LDAPSEARCH} cn="${group}" memberUid | grep -v ^dn: | awk '{print $2}')"
			if [ -n "${group_users}" ]; then
				all="${group_users} ${all}"
				verbose "${all}" 4
			fi
		fi
	fi

	RECIPIENTS="${RECIPIENTS} $(echo ${all} | sort -u)"
}

generateSessionKey() {
	verbose "Generating a session key..." 2

	openssl rand 32 -base64 -out "${SESSION_KEY}" 2>/dev/null
}

getFingerprintFromKey() {
	local readonly keyfile="${1}"
	local readonly mode="${2}"
	local fp

	verbose "Extracting fingerprint from '${keyfile}'..." 4

	fp=$(ssh-keygen -l -f "${keyfile}" | awk '{print $2}')

	# ssh-keygen does not yield errors to stderr; if the input is not
	# a valid key, it will print "mumble is not a public key" instead.

	# When decrypting, $keyfile will be a privkey, and we can only get
	# the fingerprint as above if the pubkey happens to be stored
	# alongside the privkey.  (It was worth trying.)
	#
	# Otherwise, we need to extract the pubkey from the privkey, for
	# which we need to prompt for the password.  Unfortunately we
	# can't intercept the password prompt here easily, so the user may
	# end up being prompted again later.

	if [ ! -z "${fp}" -a x"${fp}" = x"is" ]; then
		if [ x"${mode}" = x"encrypt" ]; then
			echo "${keyfile} appears not to be a valid ssh key." >&2
			return
		else
			verbose "Trying to convert '${keyfile}' to a pubkey..." 4
			ssh-keygen -y -f "${keyfile}" >"${TDIR}/pubkey" </dev/tty || {
				echo "Unable to convert ${keyfile} into a pubkey." >&2
				echo "Either the file is not a privkey or the passphrase didn't match." >&2
				return
			}
			fp=$(ssh-keygen -l -f "${TDIR}/pubkey" | awk '{print $2}')
		fi
	fi

	echo "${fp}"
}


# Retrieve all available keys for all recipients and store them in a
# single file, 'keyfile'.
getPubkeys() {
	verbose "Identifying/retrieving ssh pubkeys..." 2

	local found=0
	local etcfile group kfile tildefile
	local readonly keyfile="${TDIR}/keyfile"

	for group in ${RECIPIENT_GROUPS}; do
		expandGroup "${group}"
	done

	if [ -n "${RECIPIENTS}" ]; then
		for r in ${RECIPIENTS}; do
			found=0
			# Tilde expansion happens before variable expansion, so
			# we need to eval echo here.
			tildefile=$(eval echo ~${r}/.ssh/authorized_keys)
			etcfile="/etc/ssh/authorized_keys/${r}"
			for f in "${tildefile}" "${etcfile}"; do
				verbose "Trying '${f}'..." 3
				if [ -r "${f}" ]; then
					appendKey "${f}" "${r}"
					found=1
				fi
			done

			if [ ${found} -eq 0 ]; then
				getPubkeysFromLDAP "${r}" || true
				if [ -s "${LDAP_KEYFILE}" ]; then
					appendKey "${LDAP_KEYFILE}" "${r}"
					found=1
				fi
			fi

			if [ ${found} -eq 0 -a -n "${URLS}" ]; then
				getPubkeysFromURLs "${r}" || true
				if [ -s "${URLS_KEYFILE}" ]; then
					appendKey "${URLS_KEYFILE}" "${r}"
					found=1
				fi
			fi
		done
	fi

	if [ -n "${KEYFILES}" ]; then
		for kfile in ${KEYFILES}; do
			appendKey "${kfile}" "${kfile}"
		done
	fi

	if [ ! -s "${keyfile}" ]; then
		echo "Unable to find any usable pubkeys." >&2
		exit 1
	fi
}

getPubkeysFromLDAP() {
	local readonly user="${1}"
	local readonly oIFS="${IFS}"

	ldapVarCheck || return

	verbose "Trying to get ssh pubkeys for '${user}' from LDAP..." 3

	${LDAPSEARCH} uid="${user}" ${LDAPFIELD}				| \
		grep -v dn: 							| \
		sed -e 's/^ *//' 						| \
		tr -d '\n'							| \
		awk -v F="${LDAPFIELD}" '{ gsub(F ":* ", "\n"); print; }'	| \
		sed -e 's/^%.*//'						| \
		grep . > "${LDAP_KEYFILE}" || {
			verbose "Unable to retrieve ssh pubkey for '${user}' from LDAP."
			return
	}

	# Keys may be base64, in which case (we assume) it will not
	# contain any spaces, otherwise required for a valid ssh key.

	IFS="
"

	for line in $(cat "${LDAP_KEYFILE}"); do
		if echo "${line}" | grep -q ' '; then
			echo "${line}" >> "${LDAP_KEYFILE}.new"
		else
			echo "${line}" | eval ${BASE64}				| \
				sed -e 's/^%.*//'				| \
				grep . > "${LDAP_KEYFILE}.tmp" || {
					echo "Ignoring unusable data in LDAP for '${user}'." >&2
					continue
				}
			cat "${LDAP_KEYFILE}.tmp" >> "${LDAP_KEYFILE}.new"
			rm "${LDAP_KEYFILE}.tmp"
		fi
	done

	if [ -s "${LDAP_KEYFILE}.new" ]; then
		mv "${LDAP_KEYFILE}.new" "${LDAP_KEYFILE}"
	fi

	IFS="${oIFS}"

	if [ ! -s "${LDAP_KEYFILE}" ]; then
		echo "Unable to retrieve ssh pubkey for '${user}' from LDAP." >&2
		return
	fi
}

getPubkeysFromURLs() {
	local readonly user="${1}"
	local url

	verbose "Trying to get ssh pubkeys for '${user}' from URLs..." 3
	rm -f "${URLS_KEYFILE}"
	for url in ${URLS}; do
		url=$(echo "${url}" | sed -e "s/<user>/${user}/g")
		verbose "Trying to get '${url}'..." 4
		curl -s "${url}" >> "${URLS_KEYFILE}" || {
			echo "Unable to curl '${url}'." >&2
		}
	done
}


usage() {
	cat <<EOH
Usage: ${PROGNAME} [-EVdehv] [-f file] [-k key] [-u user]
        -E       expert mode; skip certain tests
	-V       print version information and exit
	-d       decrypt
	-e       encrypt (default)
	-f file  encrypt/decrypt this file
	-h       print this help and exit
	-k key   encrypt using this public key file
	-u user  encrypt for this user
	-v       be verbose
EOH
}

ldapVarCheck() {
	if [ -z "${LDAPSEARCH}" -o -z "${LDAPFIELD}" ]; then
		verbose "Use of LDAP not configured."
		verbose "Please set the LDAPSEARCH and LDAPFIELD variables near the top of" 2
		verbose "this script or in your environment." 2
		return 1
	fi
	return 0
}


# purpose : check that all important variables are set correctly
# input   : none
# output  : none, globals may have been modified, may abort on error

varCheck() {
	local file kfiles n

	verbose "Checking that all variables look ok..."
	n=0
	if [ -n "${KEYFILES}" ]; then
		for file in ${KEYFILES}; do
			if [ ! -r "${file}" ]; then
				echo "Skipping unreadable key file '${file}'." >&2
			else
				kfiles="${file} ${kfiles}"
				n=$(( ${n} + 1 ))
			fi
		done
		if [ -z "${kfiles}" ]; then
			echo "No usable key file found." >&2
			exit 1
		fi
		KEYFILES="${kfiles}"
	fi

	if [ x"${ACTION}" = x"encrypt" ]; then
		if [ -z "${KEYFILES}" -a -z "${RECIPIENTS}" -a -z "${RECIPIENT_GROUPS}" ]; then
			echo "You need to provide either a key file, a group, or a username." >&2
			exit 1
		fi
	else
		if [ -n "${RECIPIENT_GROUPS}" ]; then
			echo "You can't specify a group when decrypting." >&2
			exit 1
		fi
		if [ -n "${RECIPIENTS}" ]; then
			echo "You can't specify a user when decrypting." >&2
			exit 1
		fi

		if [ ${n} -gt 1 ]; then
			echo "Please only specify a single key file when decrypting." >&2
			exit 1
		fi

		if [ -z "${KEYFILES}" ]; then
			verbose "No key specified, trying ~/.ssh/id_rsa..." 2
			KEYFILES="${HOME}/.ssh/id_rsa"
			if [ ! -r "${KEYFILES}" ]; then
				echo "No key found.  Please use '-k'." >&2
				exit 1
			fi
		fi
	fi

	if [ x"${UNAME}" = x"Linux" ]; then
		BASE64="base64 -d -i"
	else
		BASE64="base64 -D"
	fi
}


verbose() {
	local msg="${1}"
	local level="${2:-1}"
	local i=0

	if [ ${level} -le ${VERBOSITY} ]; then
		while [ ${i} -lt ${level} ]; do
			printf "=" >&2
			i=$(( ${i} + 1 ))
		done
		echo "> ${msg}" >&2
	fi
}

###
### Main
###

trap 'cleanup' 0

while getopts 'EGVdef:g:hk:u:v' opt; do
	case ${opt} in
		E)
			EXPERT="yes"
		;;
		G)
			URLS="${URLS} ${GITHUB}"
		;;
		V)
			echo "jass version ${VERSION}"
			echo "Able to decrypt input generated with jass versions: ${CAN_DECRYPT_FROM}"
			echo "Able to encrypt data for jass versions: ${CAN_ENCRYPT_FOR}"
			exit 0
			# NOTREACHED
		;;
		d)
			ACTION="decrypt"
		;;
		e)
			ACTION="encrypt"
		;;
		f)
			if [ x"${OPTARG}" != x"-" ]; then
				FILE="${OPTARG}"
			fi
		;;
		g)
			RECIPIENT_GROUPS="${OPTARG} ${RECIPIENT_GROUPS}"
		;;
		h\?)
			usage
			exit 0
			# NOTREACHED
		;;
		k)
			KEYFILES="${OPTARG} ${KEYFILES}"
		;;
		u)
			RECIPIENTS="${OPTARG} ${RECIPIENTS}"
		;;
		v)
			VERBOSITY=$(( ${VERBOSITY} + 1 ))
		;;
		*)
			usage
			exit 1
			# NOTREACHED
		;;
	esac
done
shift $(($OPTIND - 1))

if [ $# -ne 0 ]; then
	usage
	exit 1
	# NOTREACHED
fi

varCheck

mkdir -p "${KEYDIR}" "${PEMDIR}"

${ACTION}

exit 0
