Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
branch: master
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

executable file 148 lines (127 sloc) 4.787 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
# utility functions for running shell commands

import re
import subprocess as spc
import threading as thg
import sys
import time

PIPE = spc.PIPE

def multi_sub(reps, s):
    """multi-substitute: do replacements (dictionary of strings to strings)
simultaneously
"""
    i = 0
    ret = ""
    for m in re.finditer("(" + "|".join(map(re.escape, reps.keys())) + ")", s):
        ret += s[i:m.start()] + reps[m.group()]
        i = m.end()
    ret += s[i:]
    return ret

# the simpler way:
# s.replace("\\", "\\\\").replace("'", "'\\''")
# does _not_ work when there are backslashes on the scene:
# you need simultaneous substitution bc \ inside ' is weird in shell-scripting
# (this really does come up, such double escaping for dsh e.g.)
arg_esc_subs = {
    "'": "'\\''",
    "\\": "'\\\\'"
}

def arg_esc(s):
    """escape an argument for a shell command"""
    for i in range(len(s)):
        if s[i] in (
                ' ', "'", '"', ';', '#', '&', '`', '\\', '!', '*', '?', '$',
                '(', '{', '[', '>', '<', '|', '~', '^'):
            return "'" + multi_sub(arg_esc_subs, s) + "'"
    return s

def args_esc(ss):
    """escape arguments for a shell command"""
    return tuple([arg_esc(s) for s in ss])

def error_out(s):
    raise Exception(s + " ..aborting!")

def cmd_run(cmd, shell=True, stdout=None, stdin=None, stderr=None, env=None):
    """
run a command, applying escaping properly on array.
basically this does what Popen should do already
(Popen is wontfix busted on unix for passing multiple args as array,
since it just does bash -c "$@" which silently kills args. pretty lame)
"""
    if type(cmd) == type([]):
        cmd = " ".join([arg_esc(a) for a in cmd])
    return spc.Popen(cmd, shell=shell, stdout=stdout, stdin=stdin,
        stderr=stderr, env=env)

def cmd_wait(cmd, shell=True, stdout=None, stderr=None, can_fail=False,
        env=None):
    """wait for a command to finish and return the subprocess object"""
    p = cmd_run(cmd, shell=shell, stdout=stdout, stderr=stderr, env=env)
    ret = p.wait()
    if not can_fail and ret != 0:
        error_out("got ret " + str(ret) + " from command:\n" + str(cmd))
    return p

def cmd_output(cmd, can_fail=False):
    """wait for a command to finish and return the output"""
    ls = cmd_wait(cmd, stdout=spc.PIPE, can_fail=can_fail).stdout.readlines()
    return [l[:-1] for l in ls]

def are_you_sure(timeout=5, msg=""):
    print("WARNING: YOU ARE PERFORMING A POTENTIALLY DANGEROUS ACTION")
    if msg:
        print()
        print(msg)
        print()
    print("waiting " + str(timeout) + " seconds before continuing.")
    print("^C TO CANCEL")
    timeout_list = range(1, timeout + 1)
    timeout_list.reverse()
    for i in timeout_list:
        print(i)
        time.sleep(1)

def try_rep(n, cmd):
    """try a command repeatedly until it works,
erroring after some number of failures
"""
    for i in xrange(n):
        if i > 0:
            print("retrying " + cmd, file=sys.stderr)
        ret = spc.Popen(cmd, shell=True).wait()
        if ret == 0:
            return ret
    print("***** FAILED ***** (with ret %d): %s" % (ret, cmd), file=sys.stderr)
    return ret

default_fanout = 5
# note this to retry the ssh connection/execution
# not to retry your command, e.g.: dsh(hosts, 'false') will only be tried
# once on each host, even though false "fails" every time it is run
default_retry_num = 10

def dsh(hosts, cmd, fanout=default_fanout, retry_num=default_retry_num):
    """runs a command on several boxes via ssh"""
    def my_cmd_fcn(x):
        return "ssh " + x + " " + cu.arg_esc(cmd)
    dsh_fcn(hosts, my_cmd_fcn, fanout, retry_num)

def dsh_fcn(hosts, cmd_fcn, fanout=default_fanout,
        retry_num=default_retry_num):
    """very generic form of dsh function"""
    print("Runinng " + cu.arg_esc(cmd_fcn("<x>")) + ":" + \
        "\n- with fanout " + str(fanout) + " and retry " + str(retry_num) + \
        "\n- on <x> in " + str(hosts))

    fanout_sema = thg.BoundedSemaphore(value=fanout)

    class ShellCommThread(thg.Thread):
        def __init__(self, cmd, retry_num):
            thg.Thread.__init__(self)
            self.cmd = cmd
            self.retry_num = retry_num
        def run(self):
            fanout_sema.acquire()
            try:
                spc.Popen(self.cmd, shell=True).wait()
            finally:
                fanout_sema.release()

    threads = []
    for host in hosts:
        t = ShellCommThread(cmd_fcn(host), retry_num)
        threads.append(t)
        t.start()
    for t in threads:
        t.join()

def sql_esc(s):
    return '"' + s.replace('"', '\"').replace('\\', '\\\\') + '"'

def sql_like_esc(s):
    return s.replace('%', '\\%').replace('_', '\\_')
Something went wrong with that request. Please try again.