def test_valid_syntax(self):
options = mitogen.parent.Options(max_message_size=123)
conn = mitogen.parent.Connection(options, self.router)
conn.context = mitogen.core.Context(None, 123)
args = conn.get_boot_command()
# Executing the boot command will print "EC0" and expect to read from
# stdin, which will fail because it's pointing at /dev/null, causing
# the forked child to crash with an EOFError and disconnect its write
# pipe. The forked and freshly execed parent will get a 0-byte read
# from the pipe, which is a valid script, and therefore exit indicating
# success.
fp = open("/dev/null", "r")
try:
proc = subprocess.Popen(
args,
stdin=fp,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = proc.communicate()
self.assertEqual(0, proc.returncode)
self.assertEqual(
stdout, mitogen.parent.BootstrapProtocol.EC0_MARKER + b('\n'))
self.assertIn(b("Error -5 while decompressing data"), stderr)
finally:
fp.close()
def test_key_shadowed_nuchange(self):
environb[b('SOMEKEY')] = b('234')
self.tf.write(b('SOMEKEY=123\n'))
self.tf.flush()
watcher = klass(self.tf.name)
watcher.check()
self.assertEqual(environb[b('SOMEKEY')], b('234'))
def test_two_grandchild_one_intermediary(self):
tf = tempfile.NamedTemporaryFile()
path = mitogen.core.to_text(tf.name)
try:
tf.write(b('test'))
tf.flush()
interm = self.router.local(name='interm')
c1 = self.router.local(via=interm, name='c1')
c2 = self.router.local(via=interm)
c1.call(prepare)
c2.call(prepare)
service = self.klass(router=self.router)
service.propagate_to(context=c1, path=path)
service.propagate_to(context=c2, path=path)
s = c1.call(wait_for_file, path=path)
self.assertEquals(b('test'), s)
s = c2.call(wait_for_file, path=path)
self.assertEquals(b('test'), s)
finally:
tf.close()
def test_no_sep(self):
left, sep, right = self.func(b('dave'), b('x'))
self.assertTrue(isinstance(left, mitogen.core.BytesType))
self.assertTrue(isinstance(sep, mitogen.core.BytesType))
self.assertTrue(isinstance(right, mitogen.core.BytesType))
self.assertEqual(left, b('dave'))
self.assertEqual(sep, b(''))
self.assertEqual(right, b(''))
def test_key_deleted(self):
environb[b('SOMEKEY')] = b('123')
self.tf.write(b('SOMEKEY=123\n'))
self.tf.flush()
watcher = klass(self.tf.name)
self.tf.seek(0)
self.tf.truncate(0)
watcher.check()
self.assertTrue(b('SOMEKEY') not in environb)
def test_main(self):
import __main__
path, src, is_pkg = self.call('__main__')
self.assertEquals(path, __main__.__file__)
# linecache adds a line ending to the final line if one is missing.
actual_src = open(path, 'rb').read()
if actual_src[-1:] != b('\n'):
actual_src += b('\n')
self.assertEquals(src, actual_src)
self.assertFalse(is_pkg)
def worker_main(self):
"""
The main function of for the mux process: setup the Mitogen broker
thread and ansible_mitogen services, then sleep waiting for the socket
connected to the parent to be closed (indicating the parent has died).
"""
self._setup_master()
self._setup_services()
# Let the parent know our listening socket is ready.
mitogen.core.io_op(self.child_sock.send, b('1'))
self.child_sock.send(b('1'))
# Block until the socket is closed, which happens on parent exit.
mitogen.core.io_op(self.child_sock.recv, 1)
def _rewrite_source(self, s):
"""
Mutate the source according to the per-task parameters.
"""
# While Ansible rewrites the #! using ansible_*_interpreter, it is
# never actually used to execute the script, instead it is a shell
# fragment consumed by shell/__init__.py::build_module_command().
new = [b('#!') + utf8(self.interpreter_fragment)]
if self.is_python:
new.append(self.b_ENCODING_STRING)
_, _, rest = bytes_partition(s, b('\n'))
new.append(rest)
return b('\n').join(new)
def _on_forward_module(self, msg):
if msg.is_dead:
return
context_id_s, _, fullname = msg.data.partition(b('\x00'))
fullname = mitogen.core.to_text(fullname)
context_id = int(context_id_s)
stream = self.router.stream_by_id(context_id)
if stream.remote_id == mitogen.parent_id:
LOG.error('%r: dropping FORWARD_MODULE(%d, %r): no route to child',
self, context_id, fullname)
return
if fullname in stream.sent_modules:
return
LOG.debug('%r._on_forward_module() sending %r to %r via %r', self,
fullname, context_id, stream.remote_id)
self._send_module_and_related(stream, fullname)
if stream.remote_id != context_id:
stream._send(
mitogen.core.Message(
data=msg.data,
handle=mitogen.core.FORWARD_MODULE,
dst_id=stream.remote_id,
))
def _parse(self, fp):
"""
linux-pam-1.3.1/modules/pam_env/pam_env.c#L207
"""
for line in fp:
# ' #export foo=some var ' -> ['#export', 'foo=some var ']
bits = shlex_split_b(line)
if (not bits) or bits[0].startswith(b('#')):
continue
if bits[0] == b('export'):
bits.pop(0)
key, sep, value = bytes_partition(b(' ').join(bits), b('='))
if key and sep:
yield key, value
def worker_main(self):
"""
The main function of the mux process: setup the Mitogen broker thread
and ansible_mitogen services, then sleep waiting for the socket
connected to the parent to be closed (indicating the parent has died).
"""
save_pid('mux')
ansible_mitogen.logging.set_process_name('mux')
ansible_mitogen.affinity.policy.assign_muxprocess(self.index)
self._setup_master()
self._setup_services()
try:
# Let the parent know our listening socket is ready.
mitogen.core.io_op(self.model.child_sock.send, b('1'))
# Block until the socket is closed, which happens on parent exit.
mitogen.core.io_op(self.model.child_sock.recv, 1)
finally:
self.broker.shutdown()
self.broker.join()
# Test frameworks living somewhere higher on the stack of the
# original parent process may try to catch sys.exit(), so do a C
# level exit instead.
os._exit(0)
def worker_main(self):
"""
The main function of the mux process: setup the Mitogen broker thread
and ansible_mitogen services, then sleep waiting for the socket
connected to the parent to be closed (indicating the parent has died).
"""
save_pid('mux')
# #623: MuxProcess ignores SIGINT because it wants to live until every
# Ansible worker process has been cleaned up by
# TaskQueueManager.cleanup(), otherwise harmles yet scary warnings
# about being unable connect to MuxProess could be printed.
signal.signal(signal.SIGINT, signal.SIG_IGN)
ansible_mitogen.logging.set_process_name('mux')
ansible_mitogen.affinity.policy.assign_muxprocess(self.index)
self._setup_master()
self._setup_services()
try:
# Let the parent know our listening socket is ready.
mitogen.core.io_op(self.model.child_sock.send, b('1'))
# Block until the socket is closed, which happens on parent exit.
mitogen.core.io_op(self.model.child_sock.recv, 1)
finally:
self.broker.shutdown()
self.broker.join()
# Test frameworks living somewhere higher on the stack of the
# original parent process may try to catch sys.exit(), so do a C
# level exit instead.
os._exit(0)
class LoadModuleTest(ImporterMixin, testlib.TestCase):
data = zlib.compress(b("data = 1\n\n"))
path = 'fake_module.py'
modname = 'fake_module'
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
response = (modname, None, path, data, [])
def test_no_such_module(self):
self.set_get_module_response(
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
(self.modname, None, None, None, None))
self.assertRaises(ImportError,
lambda: self.importer.load_module(self.modname))
def test_module_added_to_sys_modules(self):
self.set_get_module_response(self.response)
mod = self.importer.load_module(self.modname)
self.assertIs(sys.modules[self.modname], mod)
self.assertIsInstance(mod, types.ModuleType)
def test_module_file_set(self):
self.set_get_module_response(self.response)
mod = self.importer.load_module(self.modname)
self.assertEquals(mod.__file__, 'master:' + self.path)
def test_module_loader_set(self):
self.set_get_module_response(self.response)
mod = self.importer.load_module(self.modname)
self.assertIs(mod.__loader__, self.importer)
def test_module_package_unset(self):
self.set_get_module_response(self.response)
mod = self.importer.load_module(self.modname)
self.assertIsNone(mod.__package__)
def test_tuple(self):
l = (1, u'b', b('c'))
roundtrip = self.roundtrip(l)
self.assertEquals(l, roundtrip)
self.assertTrue(isinstance(roundtrip, tuple))
for k in range(len(l)):
self.assertTrue(isinstance(roundtrip[k], type(l[k])))
def test_list(self):
l = [1, u'b', b('c')]
roundtrip = self.roundtrip(l)
self.assertTrue(isinstance(roundtrip, list))
self.assertEquals(l, roundtrip)
for k in range(len(l)):
self.assertTrue(isinstance(roundtrip[k], type(l[k])))
def test_dict(self):
d = {1: 2, u'a': 3, b('b'): 4, 'c': {}}
roundtrip = self.roundtrip(d)
self.assertEquals(d, roundtrip)
self.assertTrue(isinstance(roundtrip, dict))
for k in d:
self.assertTrue(isinstance(roundtrip[k], type(d[k])))
def iter_read(fds, deadline=None):
poller = PREFERRED_POLLER()
for fd in fds:
poller.start_receive(fd)
bits = []
timeout = None
try:
while poller.readers:
if deadline is not None:
timeout = max(0, deadline - time.time())
if timeout == 0:
break
for fd in poller.poll(timeout):
s, disconnected = mitogen.core.io_op(os.read, fd, 4096)
if disconnected or not s:
IOLOG.debug('iter_read(%r) -> disconnected', fd)
poller.stop_receive(fd)
else:
IOLOG.debug('iter_read(%r) -> %r', fd, s)
bits.append(s)
yield s
finally:
poller.close()
if not poller.readers:
raise mitogen.core.StreamError(
u'EOF on stream; last 300 bytes received: %r' %
(b('').join(bits)[-300:].decode('latin1'), ))
raise mitogen.core.TimeoutError('read timed out')
def filter_debug(stream, it):
"""
Read line chunks from it, either yielding them directly, or building up and
logging individual lines if they look like SSH debug output.
This contains the mess of dealing with both line-oriented input, and partial
lines such as the password prompt.
Yields `(line, partial)` tuples, where `line` is the line, `partial` is
:data:`True` if no terminating newline character was present and no more
data exists in the read buffer. Consuming code can use this to unreliably
detect the presence of an interactive prompt.
"""
# The `partial` test is unreliable, but is only problematic when verbosity
# is enabled: it's possible for a combination of SSH banner, password
# prompt, verbose output, timing and OS buffering specifics to create a
# situation where an otherwise newline-terminated line appears to not be
# terminated, due to a partial read(). If something is broken when
# ssh_debug_level>0, this is the first place to look.
state = 'start_of_line'
buf = b('')
for chunk in it:
buf += chunk
while buf:
if state == 'start_of_line':
if len(buf) < 8:
# short read near buffer limit, block awaiting at least 8
# bytes so we can discern a debug line, or the minimum
# interesting token from above or the bootstrap
# ('password', 'MITO000\n').
break
elif buf.startswith(DEBUG_PREFIXES):
state = 'in_debug'
else:
state = 'in_plain'
elif state == 'in_debug':
if b('\n') not in buf:
break
line, _, buf = buf.partition(b('\n'))
LOG.debug('%r: %s', stream,
mitogen.core.to_text(line.rstrip()))
state = 'start_of_line'
elif state == 'in_plain':
line, nl, buf = buf.partition(b('\n'))
yield line + nl, not (nl or buf)
if nl:
state = 'start_of_line'
def filter_debug(stream, it):
"""
Read line chunks from it, either yielding them directly, or building up and
logging individual lines if they look like SSH debug output.
This contains the mess of dealing with both line-oriented input, and partial
lines such as the password prompt.
Yields `(line, partial)` tuples, where `line` is the line, `partial` is
:data:`True` if no terminating newline character was present and no more
data exists in the read buffer. Consuming code can use this to unreliably
detect the presence of an interactive prompt.
"""
# The `partial` test is unreliable, but is only problematic when verbosity
# is enabled: it's possible for a combination of SSH banner, password
# prompt, verbose output, timing and OS buffering specifics to create a
# situation where an otherwise newline-terminated line appears to not be
# terminated, due to a partial read(). If something is broken when
# ssh_debug_level>0, this is the first place to look.
state = 'start_of_line'
buf = b('')
for chunk in it:
buf += chunk
while buf:
if state == 'start_of_line':
if len(buf) < 8:
# short read near buffer limit, block awaiting at least 8
# bytes so we can discern a debug line, or the minimum
# interesting token from above or the bootstrap
# ('password', 'MITO000\n').
break
elif buf.startswith(DEBUG_PREFIXES):
state = 'in_debug'
else:
state = 'in_plain'
elif state == 'in_debug':
if b('\n') not in buf:
break
line, _, buf = buf.partition(b('\n'))
LOG.debug('%r: %s', stream, line.rstrip())
state = 'start_of_line'
elif state == 'in_plain':
line, nl, buf = buf.partition(b('\n'))
yield line + nl, not (nl or buf)
if nl:
state = 'start_of_line'
def _send_forward_module(self, stream, context, fullname):
if stream.protocol.remote_id != context.context_id:
stream.protocol._send(
mitogen.core.Message(
data=b('%s\x00%s' % (context.context_id, fullname)),
handle=mitogen.core.FORWARD_MODULE,
dst_id=stream.protocol.remote_id,
))
class JsonArgsRunner(ScriptRunner):
JSON_ARGS = b('<<INCLUDE_ANSIBLE_MODULE_JSON_ARGS>>')
def _get_args_contents(self):
return json.dumps(self.args).encode()
def _rewrite_source(self, s):
return (super(JsonArgsRunner, self)._rewrite_source(s).replace(
self.JSON_ARGS, self._get_args_contents()))
def _on_host_key_request(self, line, match):
if self.stream.conn.options.check_host_keys == 'accept':
LOG.debug('%s: accepting host key', self.stream.name)
self.stream.transmit_side.write(b('yes\n'))
return
# _host_key_prompt() should never be reached with ignore or enforce
# mode, SSH should have handled that. User's ssh_args= is conflicting
# with ours.
self.stream.conn._fail_connection(HostKeyError(hostkey_config_msg))
def _host_key_prompt(self):
if self.check_host_keys == 'accept':
LOG.debug('%r: accepting host key', self)
self.tty_stream.transmit_side.write(b('y\n'))
return
# _host_key_prompt() should never be reached with ignore or enforce
# mode, SSH should have handled that. User's ssh_args= is conflicting
# with ours.
raise HostKeyError(self.hostkey_config_msg)
def _host_key_prompt(self):
if self.check_host_keys == 'accept':
LOG.debug('%r: accepting host key', self)
self.diag_stream.transmit_side.write(b('yes\n'))
return
# _host_key_prompt() should never be reached with ignore or enforce
# mode, SSH should have handled that. User's ssh_args= is conflicting
# with ours.
raise HostKeyError(self.hostkey_config_msg)
class LoadSubmoduleTest(ImporterMixin, testlib.TestCase):
data = zlib.compress(b("data = 1\n\n"))
path = 'fake_module.py'
modname = 'mypkg.fake_module'
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
response = (modname, None, path, data, [])
def test_module_package_unset(self):
self.set_get_module_response(self.response)
mod = self.importer.load_module(self.modname)
self.assertEquals(mod.__package__, 'mypkg')
def test_stdin(self):
proc, info, _ = run_fd_check(self.func, 0, 'read',
lambda proc: proc.stdin.send(b('TEST')))
st = os.fstat(proc.stdin.fileno())
self.assertTrue(stat.S_ISSOCK(st.st_mode))
self.assertEqual(st.st_dev, info['st_dev'])
self.assertEqual(st.st_mode, _osx_mode(info['st_mode']))
flags = fcntl.fcntl(proc.stdin.fileno(), fcntl.F_GETFL)
self.assertTrue(flags & os.O_RDWR)
self.assertTrue(info['buf'], 'TEST')
self.assertTrue(info['flags'] & os.O_RDWR)
def get_boot_command(self):
source = inspect.getsource(self._first_stage)
source = textwrap.dedent('\n'.join(source.strip().split('\n')[2:]))
source = source.replace(' ', '\t')
source = source.replace('CONTEXT_NAME', self.remote_name)
preamble_compressed = self.get_preamble()
source = source.replace('PREAMBLE_COMPRESSED_LEN',
str(len(preamble_compressed)))
compressed = zlib.compress(source.encode(), 9)
encoded = codecs.encode(compressed, 'base64').replace(b('\n'), b(''))
# We can't use bytes.decode() in 3.x since it was restricted to always
# return unicode, so codecs.decode() is used instead. In 3.x
# codecs.decode() requires a bytes object. Since we must be compatible
# with 2.4 (no bytes literal), an extra .encode() either returns the
# same str (2.x) or an equivalent bytes (3.x).
return self.get_python_argv() + [
'-c',
'import codecs,os,sys;_=codecs.decode;'
'exec(_(_("%s".encode(),"base64"),"zip"))' % (encoded.decode(), )
]
def worker_main(self):
"""
The main function of for the mux process: setup the Mitogen broker
thread and ansible_mitogen services, then sleep waiting for the socket
connected to the parent to be closed (indicating the parent has died).
"""
self._setup_master()
self._setup_services()
# Let the parent know our listening socket is ready.
mitogen.core.io_op(self.child_sock.send, b('1'))
# Block until the socket is closed, which happens on parent exit.
mitogen.core.io_op(self.child_sock.recv, 1)
def neutralize_main(self, path, src):
"""Given the source for the __main__ module, try to find where it
begins conditional execution based on a "if __name__ == '__main__'"
guard, and remove any code after that point."""
match = self.MAIN_RE.search(src)
if match:
return src[:match.start()]
if b('mitogen.main(') in src:
return src
LOG.error(self.main_guard_msg, path)
raise ImportError('refused')
def test_stdin(self):
proc, info, _ = run_fd_check(self.func, 0, 'read',
lambda proc: proc.stdin.write(b('TEST')))
st = os.fstat(proc.stdin.fileno())
self.assertTrue(stat.S_ISCHR(st.st_mode))
self.assertTrue(stat.S_ISCHR(info['st_mode']))
self.assertTrue(isinstance(info['ttyname'], mitogen.core.UnicodeType))
self.assertTrue(os.isatty(proc.stdin.fileno()))
flags = fcntl.fcntl(proc.stdin.fileno(), fcntl.F_GETFL)
self.assertTrue(flags & os.O_RDWR)
self.assertTrue(info['flags'] & os.O_RDWR)
self.assertTrue(info['buf'], 'TEST')
def test_uncork(self):
context, handler = self.build()
rec = self.record()
handler.emit(rec)
handler.uncork()
self.assertEqual(1, context.send.call_count)
self.assertEqual(None, handler._buffer)
_, args, _ = context.send.mock_calls[0]
msg, = args
self.assertEqual(mitogen.core.FORWARD_LOG, msg.handle)
self.assertEqual(b('name\x0099\x00msg'), msg.data)
def exec_args(args, in_data='', chdir=None, shell=None, emulate_tty=False):
"""
Run a command in a subprocess, emulating the argument handling behaviour of
SSH.
:param list[str]:
Argument vector.
:param bytes in_data:
Optional standard input for the command.
:param bool emulate_tty:
If :data:`True`, arrange for stdout and stderr to be merged into the
stdout pipe and for LF to be translated into CRLF, emulating the
behaviour of a TTY.
:return:
(return code, stdout bytes, stderr bytes)
"""
LOG.debug('exec_args(%r, ..., chdir=%r)', args, chdir)
assert isinstance(args, list)
if emulate_tty:
stderr = subprocess.STDOUT
else:
stderr = subprocess.PIPE
proc = subprocess.Popen(
args=args,
stdout=subprocess.PIPE,
stderr=stderr,
stdin=subprocess.PIPE,
cwd=chdir,
)
stdout, stderr = proc.communicate(in_data)
if emulate_tty:
stdout = stdout.replace(b('\n'), b('\r\n'))
return proc.returncode, stdout, stderr or b('')
def test_remote_dead_message(self):
# Router should send dead message to original recipient when reply_to
# is unset.
router = self.klass(broker=self.broker, max_message_size=4096)
# Try function call. Receiver should be woken by a dead message sent by
# router due to message size exceeded.
child = router.local()
recv = mitogen.core.Receiver(router)
recv.to_sender().send(b('x') * 4097)
e = self.assertRaises(mitogen.core.ChannelError,
lambda: recv.get().unpickle())
expect = router.too_large_msg % (4096, )
self.assertEquals(e.args[0], expect)
def _connect_bootstrap(self, extra_fd):
fds = [self.receive_side.fd]
if extra_fd is not None:
self.tty_stream = mitogen.parent.DiagLogStream(extra_fd, self)
fds.append(extra_fd)
it = mitogen.parent.iter_read(fds=fds, deadline=self.connect_deadline)
password_sent = False
for buf, partial in filter_debug(self, it):
LOG.debug('%r: received %r', self, buf)
if buf.endswith(self.EC0_MARKER):
self._ec0_received()
return
elif HOSTKEY_REQ_PROMPT in buf.lower():
self._host_key_prompt()
elif HOSTKEY_FAIL in buf.lower():
raise HostKeyError(self.hostkey_failed_msg)
elif buf.lower().startswith((
PERMDENIED_PROMPT,
b("%s@%s: %s" % (self.username, self.hostname,
PERMDENIED_PROMPT)),
)):
# issue #271: work around conflict with user shell reporting
# 'permission denied' e.g. during chdir($HOME) by only matching
# it at the start of the line.
if self.password is not None and password_sent:
raise PasswordError(self.password_incorrect_msg)
elif PASSWORD_PROMPT in buf and self.password is None:
# Permission denied (password,pubkey)
raise PasswordError(self.password_required_msg)
else:
raise PasswordError(self.auth_incorrect_msg)
elif partial and PASSWORD_PROMPT in buf.lower():
if self.password is None:
raise PasswordError(self.password_required_msg)
LOG.debug('%r: sending password', self)
self.tty_stream.transmit_side.write(
(self.password + '\n').encode()
)
password_sent = True
raise mitogen.core.StreamError('bootstrap failed')
def test_valid_syntax(self):
stream = mitogen.parent.Stream(self.router, 0, max_message_size=123)
args = stream.get_boot_command()
# Executing the boot command will print "EC0" and expect to read from
# stdin, which will fail because it's pointing at /dev/null, causing
# the forked child to crash with an EOFError and disconnect its write
# pipe. The forked and freshly execed parent will get a 0-byte read
# from the pipe, which is a valid script, and therefore exit indicating
# success.
fp = open("/dev/null", "r")
proc = subprocess.Popen(args,
stdin=fp,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = proc.communicate()
self.assertEquals(0, proc.returncode)
self.assertEquals(mitogen.parent.Stream.EC0_MARKER, stdout)
self.assertIn(b("Error -5 while decompressing data: incomplete or truncated stream"), stderr)
def make(self):
return self.klass(b('x') * 128)
def test_decays_on_constructor(self):
blob = self.make()
self.assertEquals(b('x')*128, mitogen.core.BytesType(blob))
def test_empty_bytes(self):
v = mitogen.core.Blob(b(''))
self.assertEquals(b(''), roundtrip(v))
def test_decays_on_write(self):
blob = self.make()
io = BytesIO()
io.write(blob)
self.assertEquals(128, io.tell())
self.assertEquals(b('x')*128, io.getvalue())
import logging
import time
try:
from shlex import quote as shlex_quote
except ImportError:
from pipes import quote as shlex_quote
import mitogen.parent
from mitogen.core import b
LOG = logging.getLogger('mitogen')
# sshpass uses 'assword' because it doesn't lowercase the input.
PASSWORD_PROMPT = b('password')
PERMDENIED_PROMPT = b('permission denied')
HOSTKEY_REQ_PROMPT = b('are you sure you want to continue connecting (yes/no)?')
HOSTKEY_FAIL = b('host key verification failed.')
DEBUG_PREFIXES = (b('debug1:'), b('debug2:'), b('debug3:'))
def filter_debug(stream, it):
"""
Read line chunks from it, either yielding them directly, or building up and
logging individual lines if they look like SSH debug output.
This contains the mess of dealing with both line-oriented input, and partial
lines such as the password prompt.
