startx-ssh-nested

#!/bin/bash
#------------------------------------------------------------------------------
# startx-ssh-nested
#
# Remote connect to 'kitsune.underwares.org' in a nested X server (Xephyr),
# tunelling everything through ssh. It's no NX but in a pinch, it will work.
# 
# Copyright (C) 2011-2012, Alexandre Gauthier <[email protected]>
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#------------------------------------------------------------------------------


# Shell settings
set -m # enable shell job control
set -b # immediate job termination notices

# Some constants with defaults
XEPHYR="/usr/bin/Xephyr" # Xephyr binary
DNUM=3 # Display number (i.e :1), without colon.
XRES="1280x1024" # Nested server resolution.
XARGS="" # Extra Xephyr arguments (i.e. +extension name)

SSH="/usr/bin/ssh" # SSH client binary
SSHPORT=22
SSHARGS="C" #extra arguments to ssh tunnel (default is -YfC)

# Command to run on host (i.e. session manager)
SSHCMD="xterm"

USERATHOST="" # Hostname to connect to

# Status logs for components
XLOG="$(dirname "$0")/$(basename "$0")-xserver.log"
SSHLOG="$(dirname "$0")/$(basename "$0")-ssh-error.log"

# -- Function Definitions -----------------------------------------------------

#
# Return the SSH pid as integer, none on failure.
##
ssh_getpid() {
	local p="$($SSH -S $SSH_CTRL_SOCKET -O check "$USERATHOST" 2>&1)"
	if [[ $? -eq 255 ]] ; then
		echo "Unable to connect! Check $SSHLOG"
	else
		# -u is GNUism for 'unbuffered'.
		p="$(echo "$p" | sed -u -e 's/^.\+(pid=\([0-9]\+\))/\1/')"
		echo "$p"
	fi
}

#
# Cleanup leftover jobs and files before exit.
# Meant to be called by traps, but may be explicitely called.
##
run_cleanup() {
	echo "CLEANUP: "

	if [[ -S $SSH_CTRL_SOCKET ]] ; then

		# Hang onto the ssh PID, we'll need this later.
		local p="$(ssh_getpid)"

		if [[ -n "$p" ]] ; then
			printf "%s" "Sending exit signal to ssh($p):"
			
			out="$($SSH -S $SSH_CTRL_SOCKET\
				-O exit $USERATHOST 2>&1)"
			
			if [[ $? -eq 0 ]] ; then
				printf "%s" "Waiting for ssh($p) to exit..."
				wait "$p" > /dev/null 2>&1
				echo "Cleaned up SSH."
			else
				echo "Failed to communicate with ssh :("
				printf "%s" "Sending SIGKILL to pid $p..."
				kill -9 "$p" > /dev/null 2>&1
			fi
		else
			echo "WARNING: Found ssh socket but no pid?"
		fi
		
		# Socket is still there? WTF?
		# NO SILLY SOCKET, GET OUT, YOU ARE A FIFO 
		# YOU DON'T BELONG IN A FILESYSTEM!
		[[ -S $SSH_CTRL_SOCKET ]] && { rm -f $SSH_CTRL_SOCKET ; echo \
			"Removed stale socket: $SSH_CTRL_SOCKET" ; }
	fi			

	# Terminate all hanging background jobs
	jpids="$(jobs -p)"
	if [[ -n "$jpids" ]] ; then
		echo "Found leftover child processes."
		echo "Sending SIGTERM to background processes... " 
		for jpid in "$jpids" ; do
			echo -n "$jpid "
			kill "$jpid" &> /dev/null
		done
		echo "done."
	fi

	echo "Done."
}

#
# Show usage and help
##
show_usage() {
cat << EOF
Nested remote X session SSH tunnel broker
Copyleft (c) Alexandre Gauthier 2010-2011
------------------------------------------
Usage: startx-ssh-nested [-h] [OPTIONS...] user@host

  OPTIONS:
    -d num
      X11 Display number for nested server (default is '$DNUM')
    
    -g geometry
      Nested server geometry. 
      Example values are: "1024x768", "1680x1050"...
      Defaults to $XRES
  
    -p ssh_port
      TCP port to connect to on SSH server. Default is '$SSHPORT'
    
    -c remote-command
      Remote command to execute through tunnel.
      Might be your session manager, wm, or some script for instance
      'gnome-session'. Defaults to '$SSHCMD'.
    
    -h
      View usage/help (this screen)
  
  user@host
    SSH user/host combination to connect to.
    For instance, to connect as user 'knifa' to a machine named 'dominatrix',
  	  \$ $0 ... knifa@dominatrix
  
              "IT'S NOT BAD DESIGN, IT'S EVENT DRIVEN! (tm)"
           (also known as "it's not retarded it's _advanced_!")

EOF
}

# -- Argument processing ------------------------------------------------------

# Bail out if no arguments
[[ $# -lt 1 ]] &&  show_usage && exit 1

# Switches/params
while getopts ":d:g:c:hp:" opt ; do
	case "$opt" in
		"d")
			DNUM=$OPTARG
			;;
		"g")
			XRES="$OPTARG"
			;;
		"c")
			SSHCMD="$OPTARG"
			;;
		"p")
			SSHPORT="$OPTARG"
			;;
		"h")
			show_usage
			exit 1
			;;
		"?")
			echo "Unknown argument $OPTARG."
			echo "See $0 -h for help."
			exit 1
			;;
		":")
			echo "$OPTARG requires an argument."
			echo "See $0 -h for help."
			exit 1
			;;
	esac
done

# Rest of the arguments
ARGS=${@:$OPTIND}

# unnamed argument checks
if [[ ${#ARGS[@]} -eq 1 ]] ; then
	USERATHOST=$ARGS
else
	echo "Invalid number of arguments. See $0 -h for help."
	exit 1
fi

# -- Main process -------------------------------------------------------------

# Failsafe checks
[[ ! -x "$XEPHYR" ]] && { echo "Xephyr: not found: $XEPHYR" ; exit 1 ; }
[[ ! -x "$SSH" ]] && { echo "ssh: not found: $SSH" ; exit 1 ; }

# Define SSH control socket, for backgrounding
SSH_CTRL_SOCKET="$HOME/.ssh_$(echo $USERATHOST | tr '@' '-' | tr '.' '_')$$"


# Horrible traps of horrors and Bad Practices (tm)
# I am a lazy exit hook loving bastard. It's probably horrible,
# some will say I probably should use 'wait', but -- it works okay.
# Pretend it's 'event-driven' and suddendly it will look very elegant.

# Kill all jobs on shell exit, also trigger exit on INT and TERM
trap 'run_cleanup' EXIT
trap 'exit 1' SIGINT SIGTERM

# Launch Xephyr nested server as a background process
echo "Launching X Session on :${DNUM}."

($XEPHYR -br -terminate -title "X11: ${USERATHOST}" \
	-screen "${XRES}" ":${DNUM}" >> "$XLOG" 2>&1) &

echo "Waiting for X Server to settle... "
echo "(Log: $XLOG)"

# GTFO if no server was brought up :(
[[ -z "$(jobs -p)" ]] && { echo "ONOES SUDDEN XEPHYR DEATH! See $XLOG." \
	; exit 1 ; }

# Set DISPLAY to newly started server
echo "Setting $DISPLAY to newly started server..."
export DISPLAY=":${DNUM}"

# Launch tunnel
echo "Launching SSH tunnel..."

$SSH -fnY"${SSHARGS}" \
     -o Controlmaster=yes -S "${SSH_CTRL_SOCKET}" \
     "$USERATHOST" "$(printf \"%s\" \"$SSHCMD\")" 2>> "$SSHLOG"

# Planned process parenthood. 
[[ $? -eq 255 ]] && { echo "SSH Tunnel fail to launch :(" ; exit 1 ; }

echo "Waiting 2 sec. for master process to settle..."
sleep 2 

[[ -n "$(ssh_getpid)" ]] || { echo "Got no pid from ssh :(" ; exit 1 ; }
echo "Session started."

# Wait on everything and hope they will cleanly die.
# EXIT trap will cleanup if they don't.
wait &> /dev/null

echo "All child processes completed."
exit 0