'''
Copyright 2019 Google LLC

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

    https://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 platform
import queue
import select
import socket
import sys
import threading


TYPE_KEYBOARD_INPUT = 0
TYPE_TERMINAL_OUTPUT = 1
TYPE_REMOTE_CLOSED = 2
TYPE_SOCKET_TIMEOUT = 3
KEEP_ALIVE_SECONDS = 10


class ConnectionClosedError(Exception):
    def __init__(self, exit_code=None):
        self.exit_code = exit_code


class SocketTimeoutError(Exception):
    pass


def GetTtyWindowSize(fd):
    import fcntl
    import struct
    import termios
    tty_size = fcntl.ioctl(fd, termios.TIOCGWINSZ, '....')
    return struct.unpack('hh', tty_size)


class PosixConsole:
    def __init__(self, channel, inputfile):
        self.channel = channel
        self.inputfile = inputfile
        self.has_tty = False
        self.linux = os.uname()[0] == 'Linux'

    def _updateWindowSize(self, signum, stackFrame):
        if self.has_tty:
            (rows, columns) = GetTtyWindowSize(self.inputfile)
            self.channel.resize_pty(columns, rows, 0, 0)

    def _socketSendQueueLevel(self):
        # Only supported on Linux
        if not self.linux:
            return 0

        import fcntl
        import struct
        SIOCOUTQ = 0x5411
        sock = self.channel.get_transport().sock
        return struct.unpack("I", fcntl.ioctl(sock.fileno(), SIOCOUTQ, '\0\0\0\0'))[0]

    def run(self):
        import termios
        import tty
        import signal

        from termios import ISTRIP, INLCR, IGNCR, ICRNL, IXON, IXANY, IXOFF
        from termios import ISIG, ICANON, ECHO, ECHOE, ECHOK, ECHONL
        from termios import OPOST, VMIN, VTIME
        from termios import TCSADRAIN

        has_tty = False
        old_tty_attrs = None
        escape_level = 0

        try:
            old_tty_attrs = termios.tcgetattr(self.inputfile)

            (iflag, oflag, cflag, lflag, ispeed, ospeed, cc) = old_tty_attrs
            iflag &= ~(ISTRIP | INLCR | IGNCR | ICRNL | IXON | IXANY | IXOFF)
            lflag &= ~(ISIG | ICANON | ECHO | ECHOE | ECHOK | ECHONL)
            oflag &= ~OPOST
            cc[VMIN] = 1
            cc[VTIME] = 0

            newattrs = [iflag, oflag, cflag, lflag, ispeed, ospeed, cc]
            termios.tcsetattr(self.inputfile, TCSADRAIN, newattrs)

            signal.signal(signal.SIGWINCH, self._updateWindowSize)

            self.has_tty = True

        except termios.error as e:
            self.has_tty = False

        try:
            self.channel.settimeout(0)
            self.channel.get_transport().set_keepalive(KEEP_ALIVE_SECONDS)
            timeouts = 0
            tx_level = 0

            while True:
                read, write, exception = select.select([self.channel,
                                                        self.inputfile],
                                                       [], [], KEEP_ALIVE_SECONDS + 1)
                timeout = not read and not write and not exception
                if self.linux and timeout:
                    timeouts += 1
                    if timeouts == 1:
                        tx_level = self._socketSendQueueLevel()
                    else:
                        current_tx_level = self._socketSendQueueLevel()
                        if current_tx_level and current_tx_level >= tx_level:
                            raise SocketTimeoutError(socket.timeout())
                else:
                    timeouts = 0

                # data from device to host
                if self.channel in read:
                    try:
                        data = self.channel.recv(256)
                        if not data:
                            exit_code = None
                            if self.channel.exit_status_ready():
                                exit_code = self.channel.recv_exit_status()
                            raise ConnectionClosedError(exit_code=exit_code)
                        sys.stdout.write(data.decode("utf-8", errors="ignore"))
                        sys.stdout.flush()
                    except socket.timeout as e:
                        raise SocketTimeoutError(e)

                # data from host to device
                if self.inputfile in read:
                    fd = self.inputfile.fileno()
                    data = os.read(fd, 1)

                    if escape_level == 0 and data == b'\r':
                        escape_level += 1
                    elif escape_level == 1 and data == b'~':
                        escape_level += 1
                        continue
                    elif escape_level == 2 and data == b'.':
                        raise ConnectionClosedError(exit_code=0)
                    else:
                        escape_level = 0

                    if not data:
                        exit_code = None
                        if self.channel.exit_status_ready():
                            exit_code = self.channel.recv_exit_status()
                        raise ConnectionClosedError(exit_code=exit_code)

                    self.channel.send(data)
        finally:
            if self.has_tty:
                termios.tcsetattr(self.inputfile, TCSADRAIN, old_tty_attrs)


class KeyboardInputThread(threading.Thread):
    def __init__(self, queue):
        super(KeyboardInputThread, self).__init__()
        self.daemon = True
        self.queue = queue

    def run(self):
        while True:
            ch = sys.stdin.read(1)
            self.queue.put((KEYBOARD_INPUT_DATA, ch))


class TerminalOutputThread(threading.Thread):
    def __init__(self, queue, channel):
        super(TerminalOutputThread, self).__init__()
        self.daemon = True
        self.queue = queue
        self.channel = channel

    def run(self):
        while True:
            try:
                data = self.channel.recv(256)
                if not data:
                    exit_code = None
                    if self.channel.exit_status_ready():
                        exit_code = self.channel.recv_exit_status()
                    self.queue.put((TYPE_REMOTE_CLOSED, exit_code))
                    break
                data = data.decode("utf-8", errors="ignore")
                self.queue.put((TYPE_TERMINAL_OUTPUT, data))
            except socket.timeout:
                self.queue.put((TYPE_SOCKET_TIMEOUT, None))
                break

        exit_status = self.channel.recv_exit_status()
        self.queue.put((TYPE_EXIT_CODE, exit_status))


class WindowsConsole:
    def __init__(self, channel, inputfile):
        self.channel = channel
        self.dataQueue = queue.Queue()
        self.inputThread = KeyboardInputThread(self.dataQueue)
        self.outputThread = TerminalOutputThread(self.dataQueue, self.channel)

    def run(self):
        self.inputThread.start()
        self.outputThread.start()

        escape_level = 0

        while True:
            dataType, data = self.queue.get()

            if dataType == TYPE_KEYBOARD_INPUT:
                if escape_level == 0 and data == b'\r':
                    escape_level += 1
                elif escape_level == 1 and data == b'~':
                    escape_level += 1
                    continue
                elif escape_level == 2 and data == b'.':
                    raise ConnectionClosedError(exit_code=0)
                else:
                    escape_level = 0

                channel.send(data)
            if dataType == TYPE_TERMINAL_OUTPUT:
                sys.stdout.write(data)
                sys.stdout.flush()
            if dataType == TYPE_REMOTE_CLOSED:
                raise ConnectionClosedError(exit_code=data)
            if dataType == TYPE_SOCKET_TIMEOUT:
                raise SocketTimeoutError()


class Console:
    def __init__(self, channel, inputfile):
        if os.name == 'nt':
            self._console = WindowsConsole(channel, inputfile)
        else:
            self._console = PosixConsole(channel, inputfile)

    def run(self):
        return self._console.run()