Source

stagehand / utils.py

"""
stagehand.utils

Common commands and functionality for stagehand.

Requires a working fabric installation on the invoking host.

Copyright 2010-2011 Van Lindberg. Released under the BSD license.

The exported API consists of the following functions:

  report(msg): Look for whether we are running QUIET or not. If 
  not, print the provided message. 

  srun(cmd, use_sudo=False, user=None): Shortcut to run/sudo 
  commands. If use_sudo is True or user is set, then automatically
  invoke sudo... unless the user is already root.
  
  test_if(func, *args, **kwargs): Run a function and return True if 
  it succeeds, False if it fails. This is good for testing the
  environment on remote hosts.
  
  which(binary): Find out if an executable with the provided name 
  is on the $PATH of the remote host. Return either an empty string 
  (not found) or the path to the executable. Uses "which" on the 
  remote host. Exposed as a task.
  
  mktempdir(): Use remote system facilities to make a temporary 
  directory on the remote system. Return the name of the directory.
  Exposed as a task.
  
  exists(fname): See if the named file or directory exists in the 
  current working directory on the remote side. Return True or 
  False.
  
  download_remote(url, target_filename=None): Download the file 
  specified by the url to the remote filesystem. Tries using curl 
  first, followed by wget, followed by remote Python. If the target 
  filename is not specified, the final element of the url will be 
  used as a filename. Exposed as a task.
  
  update_sudo_timestamp(): Do something innocuous via sudo 
  (currently run 'ls') to force caching of the sudo password (for 
  first time sudoing) and otherwise update the sudo timestamp for 
  other uses.

  line_in_file(line, filename, quiet=True, use_sudo=False): Use a 
  remote grep to see if a particular line is in a file. Return
  True or False. Exposed as a task.
  
  is_ipaddr(addr): Test to see if the address looks like a valid IP 
  address or not. Return True or False.
  
  test_hostname(hostname): Test to see if we can connect via hostname.
  Set the .hosts attr first. Return True or False.
  
  is_dir(remote_path): Test to see if the specified remote path is
  a directory or not. Return True or False.
  
"""
# stdlib imports
from urllib import urlopen, urlretrieve
from urlparse import urlsplit
from posixpath import split as pathsplit
from contextlib import contextmanager
from functools import partial
import socket

# Third party imports
from fabric.api import *
import fabric.contrib.files as remote

# To avoid circular dependencies, do not import anything from the 
# project except .constants (which itself only imports from 
# the stdlib)

from .constants import LOCAL_BIN_DIR, LOCAL_EXT, WIN32, QUIET

__all__ = ['report', 'srun', 'test_if', 'which', 'mktempdir',
           'exists', 'download_remote', 'update_sudo_timestamp',
           'test_hostname', 'is_dir', 'is_ipaddr']

def report(msg):
  """
  Look for whether we are running QUIET or not. If not, print
  the provided message.
  """
  if not QUIET: puts(msg)


def srun(cmd, use_sudo=False, user=None):
  """
  Shortcut to run/sudo commands.
  """
  runner = sudo if (use_sudo or user) else run
  if user: 
    output = runner(cmd, user=user)
  else: 
    output = runner(cmd)
  return output
  

def test_if(func, *args, **kwargs):
  """
  Run a function and return True if it succeeds,
  False if it fails. This is good for testing the
  environment on remote hosts.
  """
  condition = False
  with settings(hide('everything'), warn_only=True):
    output = func(*args, **kwargs)
    if not output.return_code: condition = True
  return condition


@task
def which(binary):
  """
  Find out if an executable with the provided name is on the $PATH 
  of the remote host. Return either an empty string (not found) or 
  the path to the executable. Uses "which" on the remote host.
  """
  which_cmd = "which {}".format(binary)
  puts(which_cmd)
  with hide('running', 'stdout'):
    if test_if(run, which_cmd):
      return srun(which_cmd)
    else: return ""


@task
def mktempdir():
  """
  Use remote system facilities to make a temporary directory.
  Return the name of the directory.
  """
  # Try mktemp
  if which('mktemp'):
    return srun('mktemp -d')
  # Try Python
  if which('python'):
    mktmpcmd = 'python -c "import tempfile as t; print t.mkdtemp()"'
    return srun(mktmpcmd)


def exists(fname):
  """
  See if the named file or directory exists in the current working
  directory on the remote side. Return True or False.
  """
  return test_if(run, 'ls {}'.format(fname))


@task
def download_remote(url, target_filename=None):
  """
  Download the file specified by the url to the remote filesystem.
  Tries using curl first, followed by wget, followed by remote Python.
  
  If the target filename is not specified, the final element of the
  url will be used as a filename.
  """
  if not target_filename:
    target_filename = pathsplit(urlsplit(url).path)[-1]
  has_curl = test_if(lambda: run('curl --version'))
  if has_curl:
    run('curl %s -o %s' % (url, target_filename))
    return
  has_wget = test_if(lambda: run('wget --version'))
  if has_wget:
    run('wget %s -o %s' % (url, target_filename))
    return
  has_python = test_if(lambda: run('python --version'))
  if has_python:
    command = "from urllib import urlretrieve as get;get('%s','%s')" % \
              (url, target_filename)
    run('python -c "%s"' % command)
    return


def update_sudo_timestamp():
  """
  Do something innocuous via sudo to force caching of the sudo 
  password (for first time sudoing) and otherwise update the 
  timestamp for other uses.
  """
  with settings(hide('everything'), warn_only=True):
    sudo('ls /')


@task
def line_in_file(line, filename, quiet=True, use_sudo=False):
  """
  Use a remote grep to see if a particular line is in a file.
  Return True or False.
  """
  runner = srun
  grep_in_file = lambda: runner('grep "%s" "%s"' % (line, filename))
  result = test_if(grep_in_file)
  if result:
    report('Value already in file:\n{}'.format(line))
  return result


def test_hostname():
  """
  Test to see if we can connect via hostname. Set the .hosts attr 
  first. Return True or False.
  """
  return test_if(partial(run, 'ls'))


def is_dir(remote_path):
  """
  Test to see if the given remote path is a directory.Return True or 
  False.
  """
  rp = remote_path
  return (remote.exists(rp) and remote.exists('{}/.'.format(rp)))


def is_ipaddr(addr):
  """
  Test to see if the address looks like a valid IP address or not.
  Return True or False.
  """
  try: 
    socket.inet_aton(addr)
    return True
  except socket.error:
    return False


#~