django-orchestra-test/orchestra/utils/system.py

import errno
import fcntl
import getpass
import os
import re
import select
import subprocess
import sys

from django.core.management.base import CommandError


def check_root(func):
    """ Function decorator that checks if user has root permissions """
    def wrapped(*args, **kwargs):
        if getpass.getuser() != 'root':
            cmd_name = func.__module__.split('.')[-1]
            msg = "Sorry, '%s' must be executed as a superuser (root)"
            raise CommandError(msg % cmd_name)
        return func(*args, **kwargs)
    return wrapped


class _Attribute(object):
    """ Simple string subclass to allow arbitrary attribute access. """
    def __init__(self, stdout):
        self.stdout = stdout


def make_async(fd):
    """ Helper function to add the O_NONBLOCK flag to a file descriptor """
    fcntl.fcntl(fd, fcntl.F_SETFL, fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK)


def read_async(fd):
    """
    Helper function to read some data from a file descriptor, ignoring EAGAIN errors
    """
    try:
        return fd.read()
    except IOError as e:
        if e.errno != errno.EAGAIN:
            raise e
        else:
            return u''


def runiterator(command, display=False, error_codes=[0], silent=False, stdin='', force_unicode=True):
    """ Subprocess wrapper for running commands concurrently """
    if display:
        sys.stderr.write("\n\033[1m $ %s\033[0m\n" % command)
    
    p = subprocess.Popen(command, shell=True, executable='/bin/bash',
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
    
    p.stdin.write(stdin)
    p.stdin.close()
    yield
    
    make_async(p.stdout)
    make_async(p.stderr)
    
    # Async reading of stdout and sterr
    # TODO cleanup
    while True:
        # TODO https://github.com/isagalaev/ijson/issues/15
        stdout = unicode() if force_unicode else ''
        sdterr = unicode() if force_unicode else ''
        # Get complete unicode chunks
        while True:
            select.select([p.stdout, p.stderr], [], [])
            stdoutPiece = read_async(p.stdout)
            stderrPiece = read_async(p.stderr)
            try:
                stdout += unicode(stdoutPiece.decode("utf8")) if force_unicode else stdoutPiece
                sdterr += unicode(stderrPiece.decode("utf8")) if force_unicode else stderrPiece
            except UnicodeDecodeError as e:
                pass
            else:
                break
        if display and stdout:
            sys.stdout.write(stdout)
        if display and stderr:
            sys.stderr.write(stderr)
        
        state = _Attribute(stdout)
        state.stderr = sdterr
        state.return_code =  p.poll()
        yield state
        
        if state.return_code != None:
            p.stdout.close()
            p.stderr.close()
            raise StopIteration


def run(command, display=False, error_codes=[0], silent=False, stdin='', async=False, force_unicode=True):
    iterator = runiterator(command, display, error_codes, silent, stdin, force_unicode)
    iterator.next()
    if async:
        return iterator
    
    stdout = ''
    stderr = ''
    for state in iterator:
        stdout += state.stdout
        stderr += state.stderr
    
    return_code = state.return_code
    
    out = _Attribute(stdout.strip())
    err = stderr.strip()
    
    out.failed = False
    out.return_code = return_code
    out.stderr = err
    if return_code not in error_codes:
        out.failed = True
        msg = "\nrun() encountered an error (return code %s) while executing '%s'\n"
        msg = msg % (return_code, command)
        if display:
            sys.stderr.write("\n\033[1;31mCommandError: %s %s\033[m\n" % (msg, err))
        if not silent:
            raise CommandError("%s %s %s" % (msg, err, out))
    
    out.succeeded = not out.failed
    return out


def sshrun(addr, command, *args, **kwargs):
    command = command.replace("'", """'"'"'""")
    cmd = "ssh -o stricthostkeychecking=no -C root@%s '%s'" % (addr, command)
    return run(cmd, *args, **kwargs)


def get_default_celeryd_username():
    """ Introspect celeryd defaults file in order to get its username """
    user = None
    try:
        with open('/etc/default/celeryd') as celeryd_defaults:
            for line in celeryd_defaults.readlines():
                if 'CELERYD_USER=' in line:
                    user = re.findall('"([^"]*)"', line)[0]
    finally:
        if user is None:
            raise CommandError("Can not find the default celeryd username")
    return user