PageRenderTime 160ms CodeModel.GetById 50ms app.highlight 43ms RepoModel.GetById 64ms app.codeStats 1ms

/silversupport/shell.py

https://bitbucket.org/ianb/silverlining/
Python | 151 lines | 125 code | 6 blank | 20 comment | 10 complexity | b6eb5f76511e338ae3d9f40b250b20bd MD5 | raw file
  1"""Simple shell routines"""
  2
  3import os
  4import re
  5import subprocess
  6
  7__all__ = ['ssh', 'run', 'apt_install', 'shell_escape',
  8           'conditional_shell_escape', 'overwrite_file',
  9           'add_to_file']
 10
 11
 12def ssh(user, host, command, **kw):
 13    """Runs ssh to the given user/host, and runs the command.
 14
 15    Any extra keyword arguments are passed to ``run``.
 16
 17    This will use sudo for some users, and ssh in directory in other cases.
 18    """
 19    if user == 'www-data':
 20        # This is a bit tricky:
 21        user = 'www-mgr'
 22        if isinstance(command, (list, tuple)):
 23            command = ' '.join(conditional_shell_escape(i) for i in command)
 24        else:
 25            command = conditional_shell_escape(command)
 26        command = 'sudo -H -u www-data %s' % command
 27    elif isinstance(command, (list, tuple)):
 28        command = ' '.join(conditional_shell_escape(i) for i in command)
 29    ssh_args = list(kw.pop('ssh_args', []))
 30    strict_host_key_checking = kw.pop('strict_host_key_checking', False)
 31    if not strict_host_key_checking:
 32        ssh_args.extend(['-o', 'StrictHostKeyChecking=no'])
 33    return run(['ssh'] + ssh_args + ['-l', user, host, command], **kw)
 34
 35
 36def run(command, extra_env=None, stdin=None,
 37        capture_stdout=False, capture_stderr=False,
 38        cwd=None, fail_on_error=False):
 39    """Runs a command.
 40
 41    The command may be a string or (preferred) a list of
 42    command+arguments.  extra_env is extra environmental keys to set.
 43    stdin is a string to pass in.  You can capture stdout and/or
 44    stderr; if you do not they are streamed directly.  Similarly if
 45    you do not give stdin, then the command is attached to the regular
 46    stdin.
 47
 48    This always return ``(stdout, stderr, returncode)``.  The first
 49    two may be None if you didn't capture them.
 50    """
 51    env = os.environ.copy()
 52    env['LANG'] = 'C'
 53    if extra_env:
 54        env.update(extra_env)
 55    kw = {}
 56    if stdin:
 57        kw['stdin'] = subprocess.PIPE
 58    if capture_stdout:
 59        kw['stdout'] = subprocess.PIPE
 60    if capture_stderr:
 61        kw['stderr'] = subprocess.PIPE
 62    if cwd:
 63        kw['cwd'] = cwd
 64    for key, value in env.items():
 65        assert isinstance(value, str), "bad value %r: %r" % (key, value)
 66    for v in command:
 67        assert isinstance(v, str), "bad command argument %r" % v
 68    try:
 69        proc = subprocess.Popen(
 70            command, env=env, **kw)
 71    except OSError, e:
 72        raise OSError('Error while running command "%s": %s' % (
 73            ' '.join(conditional_shell_escape(i) for i in command), e))
 74    stdout, stderr = proc.communicate(stdin or '')
 75    if fail_on_error and proc.returncode:
 76        if stderr:
 77            msg = '; stderr:\n%s' % stderr
 78        else:
 79            msg = ''
 80        raise OSError(
 81            'Error while running command "%s": returncode %s%s'
 82            % (' '.join(conditional_shell_escape(i) for i in command),
 83               proc.returncode, msg))
 84    return stdout, stderr, proc.returncode
 85
 86
 87def apt_install(packages, **kw):
 88    """Install the given list of packages"""
 89    return run(['apt-get', 'install', '-q=2', '-y', '--force-yes',
 90                '--no-install-recommends'] + packages, **kw)
 91
 92
 93_end_quote_re = re.compile(r"^('*)(.*?)('*)$", re.S)
 94_quote_re = re.compile("'+")
 95
 96
 97def shell_escape(arg):
 98    """Escapes an argument for the shell"""
 99    end_quotes_match = _end_quote_re.match(arg)
100    inner = end_quotes_match.group(2)
101    inner = _quote_re.sub(lambda m: "'%s'" % m.group(0).replace("'", "\\'"),
102                          inner)
103    return ("'" + end_quotes_match.group(1).replace("'", "\\'")
104            + inner + end_quotes_match.group(3).replace("'", "\\'") + "'")
105
106
107def conditional_shell_escape(arg):
108    """Escapes an argument, unless it's obvious it doesn't need
109    escaping"""
110    if re.match(r'^[a-zA-Z0-9_,./-]+$', arg):
111        return arg
112    else:
113        return shell_escape(arg)
114
115
116def overwrite_file(dest, source_text=None, source_file=None):
117    """Overwrite a file on the system"""
118    if source_file:
119        assert not source_text
120        fp = open(source_file, 'rb')
121        source_text = fp.read()
122        fp.close()
123    if os.path.exists(dest):
124        fp = open(dest, 'rb')
125        cur_text = fp.read()
126        fp.close()
127        if cur_text == source_text:
128            return
129        new = open(dest+'.silver-back', 'wb')
130        new.write(cur_text)
131        new.close()
132    ## FIXME: is writing over the file sufficient to keep its permissions?
133    fp = open(dest, 'wb')
134    fp.write(source_text)
135    fp.close()
136
137
138def add_to_file(dest, source_text, add_newline=True):
139    """Write the text to the end of the dest, or do nothing if the text
140    is found in the file already (anywhere in the file)"""
141    fp = open(dest, 'rb')
142    content = fp.read()
143    fp.close()
144    if source_text in content:
145        return
146    if not content.endswith('\n') and add_newline:
147        content += '\n'
148    content += source_text
149    fp = open(dest, 'wb')
150    fp.write(content)
151    fp.close()