Newer
Older
tsshbatch / tsshbatch.py
#!/usr/bin/env python
# tsshbatch.py - Non-Interactive ssh Connection
# Copyright (c) 2011-2013 TundraWare Inc.
# Permission Hereby Granted For Unrestricted Personal Or Commercial Use
# See "tsshbatch-license.txt" For Licensing Details
#
# For Updates See:  http://www.tundraware.com/Software/tsshbatch

# A tip of the hat for some of the ideas in the program goes to:
#
#     http://jessenoller.com/2009/02/05/ssh-programming-with-paramiko-completely-different/

#####
# Program Housekeeping
#####

PROGNAME     = "tsshbatch.py"
BASENAME     = PROGNAME.split(".py")[0]
PROGENV      = BASENAME.upper()
CVSID        = "$Id: tsshbatch.py,v 1.158 2013/10/24 02:29:17 tundra Exp $"
VERSION      = CVSID.split()[2]

CPRT         = "(c)"
DATE         = "2011-2013"
OWNER        = "TundraWare Inc."
RIGHTS       = "All Rights Reserved."
COPYRIGHT    = "Copyright %s %s, %s  %s" % (CPRT, DATE, OWNER, RIGHTS)

PROGVER      = PROGNAME + " " + VERSION + (" - %s" % COPYRIGHT)
HOMEPAGE     = "http://www.tundraware.com/Software/%s\n" % BASENAME


#####
# Suppress Deprecation Warnings 
# Required in some older environments where paramiko version
# is behind the python libs version.
#####

import warnings
warnings.filterwarnings("ignore", "", DeprecationWarning)


#####
# Imports
#####

import getopt
import getpass
import os
import paramiko
import shlex
import socket
import sys


#####
# Constants And Literals
#####


ABORTING    = "Aborting ..."
COMMENT     = "#"
CONSUCCESS  = "SUCCESS: Connection Established"
FAILURE     = "FAILURE: %s"
HOSTSEP     = ":"
INCLUDE     = ".include"
INDENTWIDTH = 8
OPTIONSLIST = "G:H:NP:Sef:hkn:p:v"
PADWIDTH    = 12
PATHSEP     = os.sep
SEPARATOR   = " --->  "
STDIN       = "-"
SUDO        = 'sudo'
SUDOPROMPT  = 'READINGSUDOPW'
SUDOARGS    = '-S -p %s' % SUDOPROMPT
SUDOPWHINT  = '(Default: login password): '
TRAILER     = ": "
USAGE       = \
    PROGVER  + "\n"                                                                      +\
    HOMEPAGE + "\n\n"                                                                    +\
    "Usage:  tsshbatch.py [-NSehkv -G 'file dest' -P 'file dest' -f cmdfile -n name -p pw ] -H 'host ..' | hostlistfile [command arg ... ]\n" +\
    "          where,\n"                                                                 +\
    "\n"                                                                                 +\
    "            -G 'file dest'  GET file on host and write local dest directory\n"      +\
    "            -H '...'        List of targeted hosts passed as a single argument\n"   +\
    "            -N              Force prompting for username\n"                         +\
    "            -P 'file dest'  PUT local file to host dest directory\n"                +\
    "            -S              Force prompting for sudo password\n"                    +\
    "            -e              Don't report remote host stderr output\n"               +\
    "            -f cmdfile      Read commands from file\n"                              +\
    "            -h              Display help\n"                                         +\
    "            -k              Use key exchange-based authentication\n"                +\
    "            -n name         Specify login name\n"                                   +\
    "            -p pw           Specify login password\n"                               +\
    "            -v              Display extended program version information\n"


#####
# Error Messages
#####

eBADARG       =  "Invalid command line: %s!"
eBADFILE      =  "Cannot open '%s'!"
eBADSUDO      =  "sudo Failed (Check Password Or Command!) sudo Error Report:  %s"
eBADTXRQ      =  "Bad Transfer Request: %s  Must Have Exactly 1 Source And 1 Destination!"
eFXERROR      =  "File Transfer Error: %s"
eINCLUDECYCLE =  "Circular Include At: %s"
eNOCONNECT    =  "Cannot Connect To %s"
eNOHOSTS      =  "No Hosts Specified!"
eNOLOGIN      =  "Cannot Login! (Login/Password Bad?)"


#####
# Informational Messages
#####

iTXFILE = "Writing %s To %s ..."


#####
# Prompts
#####

pPASS = "Password: "
pSUDO = "%s Password: " % SUDO
pUSER = "Username (%s): "


#####
# Options That Can Be Overriden By User
####

GETSUDOPW      = False     # Prompt for sudo password
HOSTS          = []        # List of hosts to target
KEYEXCHANGE    = False     # Do key exchange-based auth?
PROMPTUSERNAME = False     # Don't use $USER, prompt for username
PWORD          = ""        # Password
REPORTERR      = True      # Report stderr output from remote host
UNAME          = ""        # Login name


#####
# Global Data Structures
#####

Commands          = []
FileIncludeStack  = []
Get_Transfer_List = {}
Put_Transfer_List = {}


#####
# Functions
#####

# Gets rid of comments and strips leading/trailing whitespace

def ConditionLine(line):
    return line.split(COMMENT)[0].strip()

# End of 'ConditionLine()'


#####
# Print Message(s) To stderr
#####

def PrintStderr(msg, TERMINATOR="\n"):
    sys.stderr.write(msg + TERMINATOR)
    sys.stderr.flush()

# End of 'PrintStderr()'


#####
# Print Message(s) To stdout
#####

def PrintStdout(msg, TERMINATOR="\n"):
    sys.stdout.write(msg + TERMINATOR)
    sys.stdout.flush()

# End of 'PrintStdout()'


#####
# Display An Error Message And Exit
#####

def ErrorExit(msg):

    PrintStderr(msg)
    os._exit(1)

# End Of 'ErrorExit()'


#####
# Transfer Files To A Host
#####

def HostFileTransfer(host, user, pw, filelist, GET=False):

    try:
        ssh = paramiko.SSHClient()

    # Connect and run the command, reporting results as we go
    
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if KEYEXCHANGE:
            ssh.connect(host)
        else:
            ssh.connect(host, username=user, password=pw)

        sftp = ssh.open_sftp()

        for srcfile in  filelist:
            for destdir in filelist[srcfile]:

                # Make sure we have a trailing path separator
                destination = destdir
                if destination[-1] != PATHSEP:
                    destination += PATHSEP

                if GET:
                    destination += host + HOSTSEP + os.path.basename(srcfile)
                    PrintStdout(iTXFILE %  (host + ":" + srcfile, destination))
                    sftp.get(srcfile, destination)

                else:
                    destination += os.path.basename(srcfile)
                    PrintStdout(iTXFILE %  (srcfile, host + ":" + destination))
                    sftp.put(srcfile, destination)

        sftp.close()
        ssh.close()

    except:

        sftp.close()
        ssh.close()
        ErrorExit(eFXERROR % (sys.exc_info()[1]))

# End of 'HostFileTransfer()'

    
def HostCommands(host, user, pw, sudopw, commands):

    ssh = paramiko.SSHClient()

    # Connect and run the command, reporting results as we go
    
    try: 
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if KEYEXCHANGE:
            ssh.connect(host)

        else:
            ssh.connect(host, username=user, password=pw)

        PrintReport([host, CONSUCCESS])

        # Run all requested commands

        for command in commands:

            # It's possible to get blank lines from stdin.
            # Ignore them.

            if not command:
                continue

            # If this is a sudo run, force password to be read
            # from stdin thereby avoiding fiddling around with ptys.

            if command.startswith(SUDO + " "):
                command = command.replace(SUDO, "%s %s" % (SUDO, SUDOARGS), 1)
          
            stdin, stdout, stderr = ssh.exec_command(command)

            # If doing a sudo command, send the password

            if command.startswith(SUDO + " "):
                stdin.write("%s\n" % sudopw)
                stdin.flush()
    
                # If all we see on stderr at this point is our original
                # prompt, then then the sudo promotion worked.  A bad
                # password or bad command will generate additional noise
                # from sudo telling us to try again or that there was a
                # command error.

                sudonoise = " ".join(stderr.readline().split(SUDOPROMPT)).strip()

                if sudonoise:                 # sudo had problems
                    PrintReport([host + " [%s]" % command, FAILURE % (eBADSUDO % sudonoise)] + ["\n"], HANDLER=PrintStderr)
                    ssh.close()
                    raise SystemExit
                    
            PrintReport([host + " (stdout)" + " [%s]" % command, "\n"] + stdout.readlines() + ["\n"])

            if REPORTERR:
                PrintReport([host + " (stderr)" + " [%s]" % command, "\n"] + stderr.readlines() + ["\n"], HANDLER=PrintStderr)

    # Handle aborts
                
    except SystemExit:
        ErrorExit(ABORTING)

    # Catch authentication problems explicitly
        
    except paramiko.AuthenticationException:
        PrintReport([host, FAILURE % eNOLOGIN], HANDLER=PrintStderr)
    
    # Everything else is some kind of connection problem

    except:
        PrintReport([host, FAILURE % (eNOCONNECT % sys.exc_info()[1].strerror)], HANDLER=PrintStderr)

    ssh.close()

# End of 'HostCommands()'


#####
# Print Report
#####

# Expects input as [host, success/failure message, result1, result2, ...]
# Uses print handler to stdout by default but can be overriden at call
# time to invoke any arbitrary handler function.

def PrintReport(results, HANDLER=PrintStdout):

    HANDLER(SEPARATOR + results[0] +
            TRAILER +
            (PADWIDTH - len(results[0])) * " " +
            results[1])

    for r in results[2:]:                             # Command Results
        HANDLER(INDENTWIDTH * " " + r.strip())

# End of 'PrintReport()'


#####
# Process A File Transfer Request
#####

def ProcessTXRQ(request, storage):

    src_dest = request.split()
    if len(src_dest) != 2:
        ErrorExit(eBADTXRQ % src_dest)

    else:

        if src_dest[0] not in storage:
            storage[src_dest[0]] = [src_dest[1],]

        else:
            storage[src_dest[0]].append(src_dest[1])

# End of 'ProcessTXRQ'

#####
# Read File Handling Comments And Directives
#####

def ReadFile(filename, listcontainer, containingfile=""):

    # Make sure we don't have a cyclic include reference

    filename = os.path.realpath(filename)
    if filename in FileIncludeStack:
        ErrorExit(eINCLUDECYCLE % containingfile + SEPARATOR + filename)

    else:
        FileIncludeStack.append(filename)  # Push it on to the stack history

    try:
        f = open(filename)
        for line in f.readlines():
            line = ConditionLine(line)

            if line:
                if line.startswith(INCLUDE):
                    fname = ConditionLine(line.split(INCLUDE)[1])
                    ReadFile(fname, listcontainer, containingfile=filename)

                else:
                    listcontainer.append(line)
        f.close()

        FileIncludeStack.pop()   # Remove this invocation from the stack
        return listcontainer

    except:
        ErrorExit(eBADFILE % filename)

# End of 'ReadFile()'


# ---------------------- Program Entry Point ---------------------- #

#####
# Process Any Options User Set In The Environment Or On Command Line
#####

# Handle any options set in the environment

OPTIONS = sys.argv[1:]
envopt = os.getenv(PROGENV)
if envopt:
    OPTIONS = shlex.split(envopt) + OPTIONS

# Combine them with those given on the command line
# This allows the command line to override defaults
# set in the environment

try:
    opts, args = getopt.getopt(OPTIONS, OPTIONSLIST)

except getopt.GetoptError, (errmsg, badarg):
    ErrorExit(eBADARG % errmsg)

for opt, val in opts:

    if opt == "-G":
        ProcessTXRQ(val, Get_Transfer_List)

    if opt == "-H":
        HOSTS = val.split()
        
    if opt == "-N":
        PROMPTUSERNAME = True
        
    if opt == "-P":
        ProcessTXRQ(val, Put_Transfer_List)

    if opt == "-S":
        GETSUDOPW = True

    if opt == "-e":
        REPORTERR = False
        
    if opt == "-f": 
        Commands = ReadFile(val, Commands)
    
    if opt == "-h":
        PrintStdout(USAGE)
        sys.exit()
        
    if opt == "-k":
        KEYEXCHANGE = True

    if opt == "-n":
        UNAME = val

    if opt == "-p":
        PWORD = val

    if opt == "-v":
        PrintStdout(CVSID)
        sys.exit()

#####
# Host & Command Line Command Definition Processing
#####

# Get the list of hosts if not specified on command line.
# The assumption is that the first argument is the file
# containing the list of targeted hosts and the remaining
# arguments form the command.

if not HOSTS:

    # Even if we are only doing file transfers and no command
    # is specified, we have to have at least one argument here
    # to tell us what hosts we're working on.

    if not args:
        ErrorExit(eNOHOSTS)

    HOSTS = ReadFile(args[0], HOSTS)
    command = " ".join(args[1:])
    
# If hosts were passed on the command line, all the arguments
# are understood to form the command.
    
else:
    command = " ".join(args[0:])


# Put it in a list data structure because this is what the
# HostCommands() function expects.  This is necessary to handle multi
# command input from from a file.

command = ConditionLine(command)
if command:
    Commands.append(command)

#####
# Authentication Credential Processing
#####

# Precedence of authentication credential sources:
#
#     1) Key exchange
#     2) Forced prompting for name via -N
#     3) Command Line/$TSSHBATCH env variable sets name
#     4) Name picked up from $USER  (Default behavior)

if not KEYEXCHANGE:

    # Preset commandline and/or program option variable username takes precedence

    if not UNAME:
        UNAME = os.getenv("USER")
        
    # By default, use the above as the login name and don't prompt for it
    # unless overriden on the command line with -N

    if PROMPTUSERNAME:

        current_user = UNAME
        UNAME = raw_input(pUSER %current_user)
        if not UNAME:                 # User just hit return - wants default
            UNAME = current_user

    # Preset commandline and/or program option variable password takes precedence

    if not PWORD:
        PWORD  = getpass.getpass(pPASS)

#####
# If Needed, Get sudo Password
####

# The need to prompt for a sudo password depends on a number of
# conditions:
#
# If a login password is present either via manual entry or -p, sudo
# will use that without further prompting.  (Default)
#
# The user is prompted for a sudo password under two conditions:
#
#  1) -k option was selected but no password was set with -p
#  2) -S option was selected 
#
# If the user IS prompted for a sudo password, any login password
# previously entered - either via -p or interactive entry - will be
# used as the default.  The user can hit enter to accept this or enter
# a different password.  This allows login and sudo passwords to be
# the same or different.

# Find out if we have any sudo commands

SUDOPRESENT = False
for command in Commands:
        if command.startswith(SUDO + " "):
            SUDOPRESENT = True

# Check condition 1) above.
# (Condition 2 handled during options processing).

if KEYEXCHANGE and not PWORD:
    GETSUDOPW = True

SUDOPW = PWORD
if SUDOPRESENT and GETSUDOPW:

    sudopwmsg   = pSUDO
    if PWORD:
        sudopwmsg = sudopwmsg[:-2] + " " + SUDOPWHINT

    SUDOPW = getpass.getpass(sudopwmsg)
    if PWORD and not SUDOPW:
        SUDOPW = PWORD

#####
# Do The Requested Work
#####

# Iterate over the list of hosts, executing any file transfers and
# commands.  Accomodate commenting out hosts in a list.

for host in HOSTS:

    if Get_Transfer_List:
        HostFileTransfer(host, UNAME, PWORD, Get_Transfer_List, GET=True)

    if Put_Transfer_List:
        HostFileTransfer(host, UNAME, PWORD, Put_Transfer_List, GET=False)

    if Commands:
        HostCommands(host, UNAME, PWORD, SUDOPW, Commands)