# dockerpty: io.py
#
# Copyright 2014 Chris Corbyn <chris@w3style.co.uk>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import fcntl
import errno
import struct
import select as builtin_select
import six


def set_blocking(fd, blocking=True):
    """
    Set the given file-descriptor blocking or non-blocking.

    Returns the original blocking status.
    """

    old_flag = fcntl.fcntl(fd, fcntl.F_GETFL)

    if blocking:
        new_flag = old_flag & ~ os.O_NONBLOCK
    else:
        new_flag = old_flag | os.O_NONBLOCK

    fcntl.fcntl(fd, fcntl.F_SETFL, new_flag)

    return not bool(old_flag & os.O_NONBLOCK)


def select(read_streams, write_streams, timeout=0):
    """
    Select the streams from `read_streams` that are ready for reading, and
    streams from `write_streams` ready for writing.

    Uses `select.select()` internally but only returns two lists of ready streams.
    """

    exception_streams = []

    try:
        return builtin_select.select(
            read_streams,
            write_streams,
            exception_streams,
            timeout,
        )[0:2]
    except builtin_select.error as e:
        # POSIX signals interrupt select()
        no = e.errno if six.PY3 else e[0]
        if no == errno.EINTR:
            return ([], [])
        else:
            raise e


class Stream(object):
    """
    Generic Stream class.

    This is a file-like abstraction on top of os.read() and os.write(), which
    add consistency to the reading of sockets and files alike.
    """

    """
    Recoverable IO/OS Errors.
    """
    ERRNO_RECOVERABLE = [
        errno.EINTR,
        errno.EDEADLK,
        errno.EWOULDBLOCK,
    ]

    def __init__(self, fd):
        """
        Initialize the Stream for the file descriptor `fd`.

        The `fd` object must have a `fileno()` method.
        """
        self.fd = fd
        self.buffer = b''
        self.close_requested = False
        self.closed = False

    def fileno(self):
        """
        Return the fileno() of the file descriptor.
        """

        return self.fd.fileno()

    def set_blocking(self, value):
        if hasattr(self.fd, 'setblocking'):
            self.fd.setblocking(value)
            return True
        else:
            return set_blocking(self.fd, value)

    def read(self, n=4096):
        """
        Return `n` bytes of data from the Stream, or None at end of stream.
        """

        while True:
            try:
                if hasattr(self.fd, 'recv'):
                    return self.fd.recv(n)
                return os.read(self.fd.fileno(), n)
            except EnvironmentError as e:
                if e.errno not in Stream.ERRNO_RECOVERABLE:
                    raise e


    def write(self, data):
        """
        Write `data` to the Stream. Not all data may be written right away.
        Use select to find when the stream is writeable, and call do_write()
        to flush the internal buffer.
        """

        if not data:
            return None

        self.buffer += data
        self.do_write()

        return len(data)

    def do_write(self):
        """
        Flushes as much pending data from the internal write buffer as possible.
        """
        while True:
            try:
                written = 0

                if hasattr(self.fd, 'send'):
                    written = self.fd.send(self.buffer)
                else:
                    written = os.write(self.fd.fileno(), self.buffer)

                self.buffer = self.buffer[written:]

                # try to close after writes if a close was requested
                if self.close_requested and len(self.buffer) == 0:
                    self.close()

                return written
            except EnvironmentError as e:
                if e.errno not in Stream.ERRNO_RECOVERABLE:
                    raise e

    def needs_write(self):
        """
        Returns True if the stream has data waiting to be written.
        """
        return len(self.buffer) > 0

    def close(self):
        self.close_requested = True

        # We don't close the fd immediately, as there may still be data pending
        # to write.
        if not self.closed and len(self.buffer) == 0:
            self.closed = True
            if hasattr(self.fd, 'close'):
                self.fd.close()
            else:
                os.close(self.fd.fileno())

    def __repr__(self):
        return "{cls}({fd})".format(cls=type(self).__name__, fd=self.fd)


class Demuxer(object):
    """
    Wraps a multiplexed Stream to read in data demultiplexed.

    Docker multiplexes streams together when there is no PTY attached, by
    sending an 8-byte header, followed by a chunk of data.

    The first 4 bytes of the header denote the stream from which the data came
    (i.e. 0x01 = stdout, 0x02 = stderr). Only the first byte of these initial 4
    bytes is used.

    The next 4 bytes indicate the length of the following chunk of data as an
    integer in big endian format. This much data must be consumed before the
    next 8-byte header is read.
    """

    def __init__(self, stream):
        """
        Initialize a new Demuxer reading from `stream`.
        """

        self.stream = stream
        self.remain = 0

    def fileno(self):
        """
        Returns the fileno() of the underlying Stream.

        This is useful for select() to work.
        """

        return self.stream.fileno()

    def set_blocking(self, value):
        return self.stream.set_blocking(value)

    def read(self, n=4096):
        """
        Read up to `n` bytes of data from the Stream, after demuxing.

        Less than `n` bytes of data may be returned depending on the available
        payload, but the number of bytes returned will never exceed `n`.

        Because demuxing involves scanning 8-byte headers, the actual amount of
        data read from the underlying stream may be greater than `n`.
        """

        size = self._next_packet_size(n)

        if size <= 0:
            return
        else:
            data = six.binary_type()
            while len(data) < size:
                nxt = self.stream.read(size - len(data))
                if not nxt:
                    # the stream has closed, return what data we got
                    return data
                data = data + nxt
            return data

    def write(self, data):
        """
        Delegates the the underlying Stream.
        """

        return self.stream.write(data)

    def needs_write(self):
        """
        Delegates to underlying Stream.
        """

        if hasattr(self.stream, 'needs_write'):
            return self.stream.needs_write()

        return False

    def do_write(self):
        """
        Delegates to underlying Stream.
        """

        if hasattr(self.stream, 'do_write'):
            return self.stream.do_write()

        return False

    def close(self):
        """
        Delegates to underlying Stream.
        """

        return self.stream.close()

    def _next_packet_size(self, n=0):
        size = 0

        if self.remain > 0:
            size = min(n, self.remain)
            self.remain -= size
        else:
            data = six.binary_type()
            while len(data) < 8:
                nxt = self.stream.read(8 - len(data))
                if not nxt:
                    # The stream has closed, there's nothing more to read
                    return 0
                data = data + nxt

            if data is None:
                return 0
            if len(data) == 8:
                __, actual = struct.unpack('>BxxxL', data)
                size = min(n, actual)
                self.remain = actual - size

        return size

    def __repr__(self):
        return "{cls}({stream})".format(cls=type(self).__name__,
                                        stream=self.stream)


class Pump(object):
    """
    Stream pump class.

    A Pump wraps two Streams, reading from one and and writing its data into
    the other, much like a pipe but manually managed.

    This abstraction is used to facilitate piping data between the file
    descriptors associated with the tty and those associated with a container's
    allocated pty.

    Pumps are selectable based on the 'read' end of the pipe.
    """

    def __init__(self,
                 from_stream,
                 to_stream,
                 wait_for_output=True,
                 propagate_close=True):
        """
        Initialize a Pump with a Stream to read from and another to write to.

        `wait_for_output` is a flag that says that we need to wait for EOF
        on the from_stream in order to consider this pump as "done".
        """

        self.from_stream = from_stream
        self.to_stream = to_stream
        self.eof = False
        self.wait_for_output = wait_for_output
        self.propagate_close = propagate_close

    def fileno(self):
        """
        Returns the `fileno()` of the reader end of the Pump.

        This is useful to allow Pumps to function with `select()`.
        """

        return self.from_stream.fileno()

    def set_blocking(self, value):
        return self.from_stream.set_blocking(value)

    def flush(self, n=4096):
        """
        Flush `n` bytes of data from the reader Stream to the writer Stream.

        Returns the number of bytes that were actually flushed. A return value
        of zero is not an error.

        If EOF has been reached, `None` is returned.
        """

        try:
            read = self.from_stream.read(n)

            if read is None or len(read) == 0:
                self.eof = True
                if self.propagate_close:
                    self.to_stream.close()
                return None

            return self.to_stream.write(read)
        except OSError as e:
            if e.errno != errno.EPIPE:
                raise e

    def is_done(self):
        """
        Returns True if the read stream is done (either it's returned EOF or
        the pump doesn't have wait_for_output set), and the write
        side does not have pending bytes to send.
        """

        return (not self.wait_for_output or self.eof) and \
                not (hasattr(self.to_stream, 'needs_write') and self.to_stream.needs_write())

    def __repr__(self):
        return "{cls}(from={from_stream}, to={to_stream})".format(
            cls=type(self).__name__,
            from_stream=self.from_stream,
            to_stream=self.to_stream)