Fabric script to manage ssh

fabfile.py

# fabfile to set up a VNC server and connection
#
# Requires xtightvncviewer to be installed, but that appears to be part of
# the default ubuntu install.
#
# There are two use cases:
# 
# * start vnc server on remote end, set up a ssh tunnel, and start vnc viewer
#   using the server and ssh tunnel set up.
# * set up a ssh tunnel and connect vnc viewer to existing vnc server
#
# Start with first case. It should also:
#
# * print out commands that can be pasted into a skype text chat window so
#   someone could connect using that command.
# * also print out stuff to use with this fabric script for the second person
#   connecting to the server.
# * stop the vnc server once done

import os, sys, time, re
from fabric.api import *
from fabric import utils
from fabric.contrib import files
from fabric.decorators import hosts
import subprocess

from tunnel import make_tunnel

VNC_SERVER = 'vnc.example.org'
# this line is to allow you to use a different ssh port
env.hosts = [VNC_SERVER + ':12345']
env.environment = 'vnc'


def startvnc(display=1, use_ssh=False):
    """Start vnc server, set up ssh tunnel and start vnc viewer."""
    display_number = int(display)
    use_ssh = _str_to_bool(use_ssh)

    _set_vnc_password()

    display_number = _check_for_existing_session(display_number)

    # start vnc server remotely
    run('tightvncserver -geometry 1024x768 -depth 15 :%s' % display_number)

    # and joing it
    joinvnc(display_number, use_ssh)
    # the above should block until the vnc session is closed, so we can carry
    # on and tear down the vncserver - the ssh tunnel should die by itself.
    killoldvnc(display_number)
    

def joinvnc(display=1, use_ssh=False):
    """Set up ssh tunnel and connect vnc viewer to existing vnc server."""
    use_ssh = _str_to_bool(use_ssh)
    display_number = int(display)
    display_port = display_number + 5900
    # start ssh tunnel
    if use_ssh:
        with make_tunnel(tunnel=env.host_string, remote='localhost:'+str(display_port), local_port=5901) as t:
            local('xtightvncviewer -compresslevel 5 -quality 8 localhost:%d' % display_port)
    else:
        local('xtightvncviewer -compresslevel 5 -quality 8 %s:%d' % (VNC_SERVER, display_port))


def killoldvnc(display=1):
    """Ensure any vnc server still running is killed."""
    display_number = int(display)
    run('tightvncserver -kill :%s' % display_number)


def _str_to_bool(bool_in):
    if isinstance(bool_in, bool):
        return bool_in
    tf = {'true': True, 'false': False}
    return tf[bool_in.lower()]

def _set_vnc_password():
    # set VNC password - see man vncpasswd
    vnc_dir = os.path.join('/home', env.user, '.vnc')
    vnc_passwd = os.path.join(vnc_dir, 'passwd')
    if not files.exists(vnc_dir):
        run('mkdir -p %s' % vnc_dir)
        run('chmod 700 %s' % vnc_dir)
    # set the password by copying the file into place
    if not files.exists(vnc_passwd):
        print "Setting VNC password"
        put('vnc_passwd', vnc_passwd, mode=0600)

def _check_for_existing_session(display_number):
    # check for existing X sessions
    with settings(warn_only=True):
        ls_out = run('ls -l /tmp/.X*-lock')
    # if no file found, no one else is running VNC
    if ls_out.return_code != 0:
        return display_number
    display_in_use = False
    display_numbers_in_use = []
    for item in ls_out.split('\n'):
        ls_out_bits = item.split(' ')
        m = re.match(r'/tmp/.X(\d+)-lock', ls_out_bits[-1])
        display_num_used = int(m.group(1))
        display_numbers_in_use.append(display_num_used)
        if display_num_used == display_number:
            display_in_use = True
            display_owner = ls_out_bits[2]
    if not display_in_use:
        return display_number
    for new_display_number in range(1,20):
        if not new_display_number in display_numbers_in_use:
            break
    if new_display_number == 20:
        utils.abort('No display numbers under 20 available! Check with a sys admin')
    print
    print 'Display number %d is in use by %s' % (display_number, display_owner)
    if display_number == 1:
        print 'Using a different display number will mean that the web interface will not be able to connect'
    use_new_num = prompt('Display number %d is available, would you like to use it? (y/n)' % new_display_number,
            validate=r'[yYnN]')
    if use_new_num.lower() != 'y':
        utils.abort('Not using different display number, aborting.')
    return new_display_number