django-orchestra/orchestra/utils/system.py

144 lines
4.1 KiB
Python
Raw Normal View History

2014-05-08 16:59:35 +00:00
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
2014-10-30 16:34:02 +00:00
class _AttributeUnicode(unicode):
2014-05-08 16:59:35 +00:00
""" Simple string subclass to allow arbitrary attribute access. """
@property
def stdout(self):
2014-10-30 16:34:02 +00:00
return unicode(self)
2014-05-08 16:59:35 +00:00
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, e:
if e.errno != errno.EAGAIN:
raise e
else:
2014-10-30 16:34:02 +00:00
return u''
2014-05-08 16:59:35 +00:00
2014-11-13 15:34:00 +00:00
def runiterator(command, display=False, error_codes=[0], silent=False, stdin=''):
""" Subprocess wrapper for running commands concurrently """
2014-05-08 16:59:35 +00:00
if display:
sys.stderr.write("\n\033[1m $ %s\033[0m\n" % command)
p = subprocess.Popen(command, shell=True, executable='/bin/bash',
2014-08-29 12:45:27 +00:00
stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
2014-08-29 16:13:34 +00:00
p.stdin.write(stdin)
p.stdin.close()
2014-11-13 15:34:00 +00:00
yield
2014-08-29 16:13:34 +00:00
2014-05-08 16:59:35 +00:00
make_async(p.stdout)
make_async(p.stderr)
2014-10-30 16:34:02 +00:00
stdout = unicode()
stderr = unicode()
2014-08-29 16:13:34 +00:00
2014-05-08 16:59:35 +00:00
# Async reading of stdout and sterr
while True:
# Wait for data to become available
select.select([p.stdout, p.stderr], [], [])
# Try reading some data from each
stdoutPiece = read_async(p.stdout)
stderrPiece = read_async(p.stderr)
if display and stdoutPiece:
sys.stdout.write(stdoutPiece)
if display and stderrPiece:
sys.stderr.write(stderrPiece)
2014-11-13 15:34:00 +00:00
return_code = p.poll()
state = _AttributeUnicode(stdoutPiece.decode("utf8"))
state.stderr = stderrPiece.decode("utf8")
state.return_code = return_code
yield state
2014-05-08 16:59:35 +00:00
2014-11-13 15:34:00 +00:00
if return_code != None:
p.stdout.close()
p.stderr.close()
raise StopIteration
def run(command, display=False, error_codes=[0], silent=False, stdin='', async=False):
iterator = runiterator(command, display, error_codes, silent, stdin)
iterator.next()
if async:
return iterator
stdout = ''
stderr = ''
for state in iterator:
stdout += state.stdout
stderr += state.stderr
return_code = state.return_code
2014-05-08 16:59:35 +00:00
2014-10-30 16:34:02 +00:00
out = _AttributeUnicode(stdout.strip())
2014-11-13 15:34:00 +00:00
err = stderr.strip()
2014-05-08 16:59:35 +00:00
out.failed = False
2014-11-13 15:34:00 +00:00
out.return_code = return_code
2014-05-08 16:59:35 +00:00
out.stderr = err
2014-11-13 15:34:00 +00:00
if return_code not in error_codes:
2014-05-08 16:59:35 +00:00
out.failed = True
msg = "\nrun() encountered an error (return code %s) while executing '%s'\n"
msg = msg % (p.returncode, command)
2014-10-02 15:58:27 +00:00
if display:
sys.stderr.write("\n\033[1;31mCommandError: %s %s\033[m\n" % (msg, err))
2014-05-08 16:59:35 +00:00
if not silent:
2014-10-04 13:23:04 +00:00
raise CommandError("%s %s %s" % (msg, err, out))
2014-05-08 16:59:35 +00:00
out.succeeded = not out.failed
return out
2014-10-03 14:02:11 +00:00
def sshrun(addr, command, *args, **kwargs):
2014-10-06 14:57:02 +00:00
command = command.replace("'", """'"'"'""")
cmd = "ssh -o stricthostkeychecking=no -C root@%s '%s'" % (addr, command)
2014-10-03 14:02:11 +00:00
return run(cmd, *args, **kwargs)
2014-05-08 16:59:35 +00:00
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