How to use the mitogen.core.b function in mitogen

To help you get started, we’ve selected a few mitogen examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github dw / mitogen / tests / testlib.py View on Github external
sock.settimeout(connect_timeout)
        try:
            sock.connect(addr)
        except socket.error:
            # Failed to connect. So wait then retry.
            time.sleep(sleep)
            continue

        if not pattern:
            # Success: We connected & there's no banner check to perform.
            sock.shutdown(socket.SHUTD_RDWR)
            sock.close()
            return

        sock.settimeout(receive_timeout)
        data = mitogen.core.b('')
        found = False
        while mitogen.core.now() < end:
            try:
                resp = sock.recv(1024)
            except socket.timeout:
                # Server stayed up, but had no data. Retry the recv().
                continue

            if not resp:
                # Server went away. Wait then retry the connection.
                time.sleep(sleep)
                break

            data += resp
            if re.search(mitogen.core.b(pattern), data):
                found = True
github dw / mitogen / mitogen / parent.py View on Github external
LOG.debug('%s: (unrecognized partial): %s',
            self.stream.name, line.decode('utf-8', 'replace'))


class BootstrapProtocol(RegexProtocol):
    """
    Respond to stdout of a child during bootstrap. Wait for :attr:`EC0_MARKER`
    to be written by the first stage to indicate it can receive the bootstrap,
    then await :attr:`EC1_MARKER` to indicate success, and
    :class:`MitogenProtocol` can be enabled.
    """
    #: Sentinel value emitted by the first stage to indicate it is ready to
    #: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have
    #: length of at least `max(len('password'), len('debug1:'))`
    EC0_MARKER = b('MITO000')
    EC1_MARKER = b('MITO001')
    EC2_MARKER = b('MITO002')

    def __init__(self, broker):
        super(BootstrapProtocol, self).__init__()
        self._writer = mitogen.core.BufferedWriter(broker, self)

    def on_transmit(self, broker):
        self._writer.on_transmit(broker)

    def _on_ec0_received(self, line, match):
        LOG.debug('%r: first stage started succcessfully', self)
        self._writer.write(self.stream.conn.get_preamble())

    def _on_ec1_received(self, line, match):
        LOG.debug('%r: first stage received mitogen.core source', self)
github dw / mitogen / ansible_mitogen / runner.py View on Github external
"""
        self.args_fp.close()
        super(ArgsFileRunner, self).revert()


class BinaryRunner(ArgsFileRunner, ProgramRunner):
    pass


class ScriptRunner(ProgramRunner):
    def __init__(self, interpreter_fragment, is_python, **kwargs):
        super(ScriptRunner, self).__init__(**kwargs)
        self.interpreter_fragment = interpreter_fragment
        self.is_python = is_python

    b_ENCODING_STRING = b('# -*- coding: utf-8 -*-')

    def _get_program(self):
        return self._rewrite_source(
            super(ScriptRunner, self)._get_program()
        )

    def _get_argv(self):
        return [
            self.args.get('_ansible_shell_executable', '/bin/sh'),
            '-c',
            self._get_shell_fragment(),
        ]

    def _get_shell_fragment(self):
        """
        Scripts are eligible for having their hashbang line rewritten, and to
github dw / mitogen / ansible_mitogen / runner.py View on Github external
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)
github dw / mitogen / ansible_mitogen / planner.py View on Github external
def read_file(path):
    fd = os.open(path, os.O_RDONLY)
    try:
        bits = []
        chunk = True
        while True:
            chunk = os.read(fd, 65536)
            if not chunk:
                break
            bits.append(chunk)
    finally:
        os.close(fd)

    return mitogen.core.b('').join(bits)
github dw / mitogen / mitogen / fork.py View on Github external
def _child_main(self, childfp):
        on_fork()
        if self.options.on_fork:
            self.options.on_fork()
        mitogen.core.set_block(childfp.fileno())

        childfp.send(b('MITO002\n'))

        # Expected by the ExternalContext.main().
        os.dup2(childfp.fileno(), 1)
        os.dup2(childfp.fileno(), 100)

        # Overwritten by ExternalContext.main(); we must replace the
        # parent-inherited descriptors that were closed by Side._on_fork() to
        # avoid ExternalContext.main() accidentally allocating new files over
        # the standard handles.
        os.dup2(childfp.fileno(), 0)

        # Avoid corrupting the stream on fork crash by dupping /dev/null over
        # stderr. Instead, handle_child_crash() uses /dev/tty to log errors.
        devnull = os.open('/dev/null', os.O_WRONLY)
        if devnull != 2:
            os.dup2(devnull, 2)
github dw / mitogen / ansible_mitogen / target.py View on Github external
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('')
github dw / mitogen / ansible_mitogen / runner.py View on Github external
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
github dw / mitogen / mitogen / service.py View on Github external
state.lock.acquire()
            try:
                for sender, fp in reversed(state.jobs):
                    sender.close()
                    fp.close()
                    state.jobs.pop()
            finally:
                state.lock.release()

    # The IO loop pumps 128KiB chunks. An ideal message is a multiple of this,
    # odd-sized messages waste one tiny write() per message on the trailer.
    # Therefore subtract 10 bytes pickle overhead + 24 bytes header.
    IO_SIZE = mitogen.core.CHUNK_SIZE - (mitogen.core.Message.HEADER_LEN + (
        len(
            mitogen.core.Message.pickled(
                mitogen.core.Blob(b(' ') * mitogen.core.CHUNK_SIZE)
            ).data
        ) - mitogen.core.CHUNK_SIZE
    ))

    def _schedule_pending_unlocked(self, state):
        """
        Consider the pending transfers for a stream, pumping new chunks while
        the unacknowledged byte count is below :attr:`window_size_bytes`. Must
        be called with the FileStreamState lock held.

        :param FileStreamState state:
            Stream to schedule chunks for.
        """
        while state.jobs and state.unacked < self.window_size_bytes:
            sender, fp = state.jobs[0]
            s = fp.read(self.IO_SIZE)