dotwryn/zsh/ssh

#!/bin/zsh

alias c='WRYN_REMOTE_CONNECTOR'
alias cmux='WRYN_REMOTE_CONNECTOR tmux'
alias sshc='WRYN_REMOTE_CONNECTOR ssh'
alias sshe='$EDITOR $WRYN_SSH_CONNECTION_LIST'
alias ssh-add-all='WRYN_SSH_ADD_CREDENTIALS'

WRYN_SSH_CONNECTION_LIST="$HOME/.config/wryn/ssh-connections"
WRYN_SSH_DELIMITER=' *| '

[ -d "$HOME/.config/wryn" ] && [ ! -f "$WRYN_SSH_CONNECTION_LIST" ] && touch "$WRYN_SSH_CONNECTION_LIST";

WRYN_SSH_ADD_CREDENTIALS() {
	for key in $(grep -slR "PRIVATE" "$HOME/.ssh"); do
		ssh-add "$key";
	done
}

function WRYN_REMOTE_CONNECTOR() {
	local COMMAND="$1";
	local CONNECTION="$2";

	local SSH_ID="^$CONNECTION$WRYN_SSH_DELIMITER"

	[ $2 ] \
		|| { echo 'Must specify a connection'; return 1 };
	grep -q "$SSH_ID" "$WRYN_SSH_CONNECTION_LIST" >/dev/null 2>&1 \
		|| { echo 'No such connection'; return 2 };

	case $COMMAND in;
		sshfs   ) sshfs $(WRYN_SSH_PARSE_ARGS $SSH_ID):$3 ${@:4}    ;;
		ssh     ) ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) ${@:3}         ;;
		cp-to   ) scp ${@:5} $(WRYN_SCP_PARSE_ARGS $SSH_ID $3):$4   ;;
		cp-from ) scp ${@:5} $(WRYN_SCP_PARSE_ARGS $SSH_ID):$3 $4   ;;
		copy-id ) ssh-copy-id ${@:3} $(WRYN_SSH_PARSE_ARGS $SSH_ID) ;;
		tmux    ) WRYN_REMOTE_TMUX "$SSH_ID" ${@:3}                 ;;

		* ) echo "$COMMAND not implemented" ;;
	esac
}
_WRYN_REMOTE_CONNECTOR() {
	local state;
	_arguments \
		'1: :->remote_command'\
		'2: :->connections'\
		':: :->remote_args'\
	;
	case "$state" in
		remote_command)
			compadd -X 'Choose a connection command:' \
				ssh sshfs cp-to cp-from copy-id tmux
			;;
		connections)
			compadd -X 'Choose a connection:' $(sed -n "s/$WRYN_SSH_DELIMITER.*//p" "$WRYN_SSH_CONNECTION_LIST")
			;;
		remote_args)
			_files
			case $words[2] in;
				ssh     ) _ssh;;
				sshfs   ) _sshfs;;
				cpto | cpfrom ) _scp;;
				tmux ) _tmux;;
				* ) ;;
			esac
			;;
	esac
}
compdef _WRYN_REMOTE_CONNECTOR WRYN_REMOTE_CONNECTOR;

WRYN_SSH_PARSE_ARGS() {
	SSH_ID="$1";
	sed -n "s/$SSH_ID\(.*\)$/\1/p" "$WRYN_SSH_CONNECTION_LIST"
}

WRYN_SCP_PARSE_ARGS() {
	SSH_ID="$1";
	[ $2 ]\
		&& sed -n "s>$SSH_ID *-p \([^ ]*\)\(.*\)$>-P \1 $2\2>p" "$WRYN_SSH_CONNECTION_LIST"\
		|| sed -n "s/$SSH_ID *-p\(.*\)$/-P\1/p" "$WRYN_SSH_CONNECTION_LIST"
}

WRYN_REMOTE_TMUX() {
	local SSH_ID="$1"
	local TMUX_UTF8_FLAG=$(locale charmap | grep -q 'UTF-8' && echo '-u')
	local DEFAULT_TMUX_SESSION_NAME='main'

	[ ! $2 ] \
		&& ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux new-session -A -s $DEFAULT_TMUX_SESSION_NAME;

	[ $3 ] \
		&& ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux $TMUX_UTF8_FLAG ${@:2} \
		|| ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux $TMUX_UTF8_FLAG a -t $2;
}