"""
Functionality for running arbitrary shell scripts.

@var ALL_USERS: A token indicating all users should be allowed.
"""
import os
import pwd
import tempfile
import operator
import shutil

from twisted.internet.protocol import ProcessProtocol
from twisted.internet.defer import (
Deferred, fail, inlineCallbacks, returnValue, succeed)
from twisted.internet.error import ProcessDone

from landscape import VERSION
from landscape.constants import UBUNTU_PATH
from landscape.lib.scriptcontent import build_script
from landscape.lib.fetch import fetch_async, HTTPCodeError
from landscape.lib.persist import Persist
from landscape.lib.encoding import encode_if_needed
from landscape.manager.plugin import ManagerPlugin, SUCCEEDED, FAILED


ALL_USERS = object()
TIMEOUT_RESULT = 102
PROCESS_FAILED_RESULT = 103
FETCH_ATTACHMENTS_FAILED_RESULT = 104


class UnknownUserError(Exception):
pass


def get_user_info(username=None):
uid = None
gid = None
path = None
if username is not None:
username_str = encode_if_needed(username)
try:
info = pwd.getpwnam(username_str)
except KeyError:
raise UnknownUserError(u"Unknown user '%s'" % username)
uid = info.pw_uid
gid = info.pw_gid
path = info.pw_dir
if not os.path.exists(path):
path = "/"
return (uid, gid, path)


class ProcessTimeLimitReachedError(Exception):
"""
Raised when a process has been running for too long.

@ivar data: The data that the process printed before reaching the time
limit.
"""

def __init__(self, data):
self.data = data


class ProcessFailedError(Exception):
"""Raised when a process exits with a non-0 exit code.

@ivar data: The data that the process printed before reaching the time
limit.
"""

def __init__(self, data, exit_code):
self.data = data
self.exit_code = exit_code


class UnknownInterpreterError(Exception):
"""Raised when the interpreter specified to run a script is invalid.

@ivar interpreter: the interpreter specified for the script.
"""

def __init__(self, interpreter):
self.interpreter = interpreter
Exception.__init__(self, self._get_message())

def _get_message(self):
return "Unknown interpreter: '%s'" % self.interpreter


class ScriptRunnerMixin(object):
"""
@param process_factory: The L{IReactorProcess} provider to run the
process with.
"""

def __init__(self, process_factory=None):
if process_factory is None:
from twisted.internet import reactor as process_factory
self.process_factory = process_factory

def is_user_allowed(self, user):
allowed_users = self.registry.config.get_allowed_script_users()
return allowed_users == ALL_USERS or user in allowed_users

def write_script_file(self, script_file, filename, shell, code, uid, gid):
# Chown and chmod it before we write the data in it - the script may
# have sensitive content
# It would be nice to use fchown(2) and fchmod(2), but they're not
# available in python and using it with ctypes is pretty tedious, not
# to mention we can't get errno.
os.chmod(filename, 0700)
if uid is not None:
os.chown(filename, uid, gid)
script_file.write(build_script(shell, code))
script_file.close()

def _run_script(self, filename, uid, gid, path, env, time_limit):

if uid == os.getuid():
uid = None
if gid == os.getgid():
gid = None

pp = ProcessAccumulationProtocol(
self.registry.reactor, self.size_limit)
self.process_factory.spawnProcess(
pp, filename, uid=uid, gid=gid, path=path, env=env)
if time_limit is not None:
pp.schedule_cancel(time_limit)
return pp.result_deferred


class ScriptExecutionPlugin(ManagerPlugin, ScriptRunnerMixin):
"""A plugin which allows execution of arbitrary shell scripts.

@ivar size_limit: The number of bytes at which to truncate process output.
"""

size_limit = 500000

def register(self, registry):
super(ScriptExecutionPlugin, self).register(registry)
registry.register_message(
"execute-script", self._handle_execute_script)

def _respond(self, status, data, opid, result_code=None):
if not isinstance(data, unicode):
# Let's decode result-text, replacing non-printable
# characters
data = data.decode("utf-8", "replace")
message = {"type": "operation-result",
"status": status,
"result-text": data,
"operation-id": opid}
if result_code:
message["result-code"] = result_code
return self.registry.broker.send_message(
message, self._session_id, True)

def _handle_execute_script(self, message):
opid = message["operation-id"]
try:
user = message["username"]
if not self.is_user_allowed(user):
return self._respond(
FAILED,
u"Scripts cannot be run as user %s." % (user,),
opid)
server_supplied_env = message.get("env", None)

d = self.run_script(message["interpreter"], message["code"],
time_limit=message["time-limit"], user=user,
attachments=message["attachments"],
server_supplied_env=server_supplied_env)
d.addCallback(self._respond_success, opid)
d.addErrback(self._respond_failure, opid)
return d
except Exception, e:
self._respond(FAILED, self._format_exception(e), opid)
raise

def _format_exception(self, e):
return u"%s: %s" % (e.__class__.__name__, e.args[0])

def _respond_success(self, data, opid):
return self._respond(SUCCEEDED, data, opid)

def _respond_failure(self, failure, opid):
code = None
if failure.check(ProcessTimeLimitReachedError):
code = TIMEOUT_RESULT
elif failure.check(ProcessFailedError):
code = PROCESS_FAILED_RESULT
elif failure.check(HTTPCodeError):
code = FETCH_ATTACHMENTS_FAILED_RESULT
return self._respond(
FAILED, str(failure.value), opid,
FETCH_ATTACHMENTS_FAILED_RESULT)

if code is not None:
return self._respond(FAILED, failure.value.data, opid, code)
else:
return self._respond(FAILED, str(failure), opid)

@inlineCallbacks
def _save_attachments(self, attachments, uid, gid, computer_id):
root_path = self.registry.config.url.rsplit("/", 1)[0] + "/attachment/"
attachment_dir = tempfile.mkdtemp()
headers = {"User-Agent": "landscape-client/%s" % VERSION,
"Content-Type": "application/octet-stream",
"X-Computer-ID": computer_id}
for filename, attachment_id in attachments.items():
if isinstance(attachment_id, str):
# Backward compatible behavior
data = attachment_id
yield succeed(None)
else:
data = yield fetch_async(
"%s%d" % (root_path, attachment_id),
cainfo=self.registry.config.ssl_public_key,
headers=headers)
full_filename = os.path.join(attachment_dir, filename)
attachment = file(full_filename, "wb")
os.chmod(full_filename, 0600)
if uid is not None:
os.chown(full_filename, uid, gid)
attachment.write(data)
attachment.close()
os.chmod(attachment_dir, 0700)
if uid is not None:
os.chown(attachment_dir, uid, gid)
returnValue(attachment_dir)

def run_script(self, shell, code, user=None, time_limit=None,
attachments=None, server_supplied_env=None):
"""
Run a script based on a shell and the code.

A file will be written with #!<shell> as the first line, as executable,
and run as the given user.

XXX: Handle the 'reboot' and 'killall landscape-client' commands
gracefully.

@param shell: The interpreter to use.
@param code: The code to run.
@param user: The username to run the process as.
@param time_limit: The number of seconds to allow the process to run
before killing it and failing the returned Deferred with a
L{ProcessTimeLimitReachedError}.
@param attachments: C{dict} of filename/data attached to the script.

@return: A deferred that will fire with the data printed by the process
or fail with a L{ProcessTimeLimitReachedError}.
"""
if not os.path.exists(shell.split()[0]):
return fail(
UnknownInterpreterError(shell))

uid, gid, path = get_user_info(user)

fd, filename = tempfile.mkstemp()
script_file = os.fdopen(fd, "w")
self.write_script_file(
script_file, filename, shell, code, uid, gid)

env = {"PATH": UBUNTU_PATH, "USER": user or "", "HOME": path or ""}
if server_supplied_env:
env.update(server_supplied_env)
old_umask = os.umask(0022)

if attachments:
persist = Persist(
filename=os.path.join(self.registry.config.data_path,
"broker.bpickle"))
persist = persist.root_at("registration")
computer_id = persist.get("secure-id")
d = self._save_attachments(attachments, uid, gid, computer_id)
else:
d = succeed(None)

def prepare_script(attachment_dir):

if attachment_dir is not None:
env["LANDSCAPE_ATTACHMENTS"] = attachment_dir

return self._run_script(
filename, uid, gid, path, env, time_limit)

d.addCallback(prepare_script)
return d.addBoth(self._cleanup, filename, env, old_umask)

def _cleanup(self, result, filename, env, old_umask):
try:
os.unlink(filename)
except:
pass
if "LANDSCAPE_ATTACHMENTS" in env:
try:
shutil.rmtree(env["LANDSCAPE_ATTACHMENTS"])
except:
pass
os.umask(old_umask)
return result


class ProcessAccumulationProtocol(ProcessProtocol):
"""A ProcessProtocol which accumulates output.

@ivar size_limit: The number of bytes at which to truncate output.
"""

def __init__(self, reactor, size_limit):
self.data = []
self.result_deferred = Deferred()
self._cancelled = False
self.size_limit = size_limit
self.reactor = reactor
self._scheduled_cancel = None

def schedule_cancel(self, time_limit):
self._scheduled_cancel = self.reactor.call_later(
time_limit, self._cancel)

def childDataReceived(self, fd, data):
"""Some data was received from the child.

Add it to our buffer, as long as it doesn't go over L{size_limit}
bytes.
"""
current_size = reduce(operator.add, map(len, self.data), 0)
self.data.append(data[:self.size_limit - current_size])

def processEnded(self, reason):
"""Fire back the deferred.

The deferred will be fired with the string of data received from the
subprocess, or if the subprocess was cancelled, a
L{ProcessTimeLimitReachedError} will be fired with data accumulated so
far.
"""
exit_code = reason.value.exitCode
data = "".join(self.data)
if self._cancelled:
self.result_deferred.errback(ProcessTimeLimitReachedError(data))
else:
if self._scheduled_cancel is not None:
scheduled = self._scheduled_cancel
self._scheduled_cancel = None
self.reactor.cancel_call(scheduled)

if reason.check(ProcessDone):
self.result_deferred.callback(data)
else:
self.result_deferred.errback(ProcessFailedError(data,
exit_code))

def _cancel(self):
"""
Close filedescriptors, kill the process, and indicate that a
L{ProcessTimeLimitReachedError} should be fired on the deferred.
"""
# Sometimes children of the shell we're killing won't die unless their
# file descriptors are closed! For example, if /bin/sh -c "cat" is the
# process, "cat" won't die when we kill its shell. I'm not sure if this
# is really sufficient: maybe there's a way we can walk over all
# children of the process we started and kill them all.
for i in (0, 1, 2):
self.transport.closeChildFD(i)
self.transport.signalProcess("KILL")
self._cancelled = True


class ScriptExecution(ManagerPlugin):
"""
Meta-plugin wrapping ScriptExecutionPlugin and CustomGraphPlugin.
"""

def __init__(self):
from landscape.manager.customgraph import CustomGraphPlugin
self._script_execution = ScriptExecutionPlugin()
self._custom_graph = CustomGraphPlugin()

def register(self, registry):
super(ScriptExecution, self).register(registry)
self._script_execution.register(registry)
self._custom_graph.register(registry)

def exchange(self, urgent=False):
self._custom_graph.exchange(urgent)