wexpect-venv/wexpect/console_reader.py

"""Wexpect is a Windows variant of pexpect https://pexpect.readthedocs.io.

Wexpect is a Python module for spawning child applications and controlling
them automatically.

console_reader Implements a virtual terminal, and starts the child program.
The main wexpect.spawn class connect to this class to reach the child's terminal.
"""

import time
import logging
import os
import traceback
import psutil
from io import StringIO

import ctypes
from ctypes import windll
import win32console
import win32process
import win32con
import win32file
import win32gui
import win32pipe
import socket

from .wexpect_util import init_logger
from .wexpect_util import EOF_CHAR
from .wexpect_util import SIGNAL_CHARS
from .wexpect_util import TIMEOUT

#
# System-wide constants
#
screenbufferfillchar = '\4'
maxconsoleY = 8000
default_port = 4321

#
# Create logger: We write logs only to file. Printing out logs are dangerous, because of the deep
# console manipulation.
#
logger = logging.getLogger('wexpect')


class ConsoleReaderBase:
    """Consol class (aka. client-side python class) for the child.

    This class initialize the console starts the child in it and reads the console periodically.
    """

    def __init__(self, path, host_pid, codepage=None, window_size_x=80, window_size_y=25,
                 buffer_size_x=80, buffer_size_y=16000, local_echo=True, interact=False, **kwargs):
        """Initialize the console starts the child in it and reads the console periodically.

        Args:
            path (str): Child's executable with arguments.
            parent_pid (int): Parent (aka. host) process process-ID
            codepage (:obj:, optional): Output console code page.
        """
        self.lastRead = 0
        self.__bufferY = 0
        self.lastReadData = ""
        self.totalRead = 0
        self.__buffer = StringIO()
        self.__currentReadCo = win32console.PyCOORDType(0, 0)
        self.pipe = None
        self.connection = None
        self.consin = None
        self.consout = None
        self.local_echo = local_echo
        self.console_pid = os.getpid()
        self.host_pid = host_pid
        self.host_process = psutil.Process(host_pid)
        self.child_process = None
        self.child_pid = None
        self.enable_signal_chars = True
        self.timeout = 30
        self.child_exitstatus = None

        logger.info(f'ConsoleReader started. location {os.path.abspath(__file__)}')

        if codepage is None:
            codepage = windll.kernel32.GetACP()

        try:
            logger.info("Setting console output code page to %s" % codepage)
            win32console.SetConsoleOutputCP(codepage)
            logger.info(
                "Console output code page: %s" % ctypes.windll.kernel32.GetConsoleOutputCP())
        except Exception as e: # pragma: no cover
            # I hope this code is unreachable...
            logger.error(e)

        try:
            self.create_connection(**kwargs)
            logger.info('Spawning %s' % path)
            try:
                self.initConsole()
                si = win32process.GetStartupInfo()
                self.__childProcess, _, self.child_pid, self.child_tid = win32process.CreateProcess(
                    None, path, None, None, False, 0, None, None, si)
                self.child_process = psutil.Process(self.child_pid)

                logger.info(f'Child pid: {self.child_pid}  Console pid: {self.console_pid}')

            except Exception: # pragma: no cover
                # I hope this code is unreachable...
                logger.error(traceback.format_exc())
                return

            if interact:
                self.interact()
                self.interact()

            self.read_loop()
        except Exception: # pragma: no cover
            # I hope this code is unreachable...
            logger.error(traceback.format_exc())
        finally:
            try:
                self.terminate_child()
                time.sleep(.01)
                self.send_to_host(self.readConsoleToCursor())
                self.sendeof()
                time.sleep(.1)
                self.close_connection()
                logger.info('Console finished.')
            except Exception: # pragma: no cover
                # I hope this code is unreachable...
                logger.error(traceback.format_exc())

    def read_loop(self):

        while True:
            if not self.isalive(self.host_process):
                logger.info('Host process has been died.')
                return

            self.child_exitstatus = win32process.GetExitCodeProcess(self.__childProcess)
            if self.child_exitstatus != win32con.STILL_ACTIVE:
                logger.info(f'Child finished with code: {self.child_exitstatus}')
                return

            consinfo = self.consout.GetConsoleScreenBufferInfo()
            cursorPos = consinfo['CursorPosition']

            if cursorPos.Y > maxconsoleY:
                '''If the console output becomes long, we suspend the child, read all output then
                clear the console before we resume the child.
                '''
                logger.info('cursorPos %s' % cursorPos)
                self.suspend_child()
                time.sleep(.2)
                self.send_to_host(self.readConsoleToCursor())
                self.refresh_console()
                self.resume_child()
            else:
                self.send_to_host(self.readConsoleToCursor())

            s = self.get_from_host()
            if s:
                logger.debug(f'get_from_host: {s}')
            else:
                logger.spam(f'get_from_host: {s}')
            if self.enable_signal_chars:
                for sig, char in SIGNAL_CHARS.items():
                    if char in s:
                        self.child_process.send_signal(sig)
            s = s.decode()
            self.write(s)


            time.sleep(.02)

    def suspend_child(self):
        """Pauses the main thread of the child process."""
        handle = windll.kernel32.OpenThread(win32con.THREAD_SUSPEND_RESUME, 0, self.child_tid)
        win32process.SuspendThread(handle)

    def resume_child(self):
        """Un-pauses the main thread of the child process."""
        handle = windll.kernel32.OpenThread(win32con.THREAD_SUSPEND_RESUME, 0, self.child_tid)
        win32process.ResumeThread(handle)

    def refresh_console(self):
        """Clears the console after pausing the child and
        reading all the data currently on the console."""

        orig = win32console.PyCOORDType(0, 0)
        self.consout.SetConsoleCursorPosition(orig)
        self.__currentReadCo.X = 0
        self.__currentReadCo.Y = 0
        writelen = self.__consSize.X * self.__consSize.Y
        # Use NUL as fill char because it displays as whitespace
        # (if we interact() with the child)
        self.consout.FillConsoleOutputCharacter(screenbufferfillchar, writelen, orig)

        self.__bufferY = 0
        self.__buffer.truncate(0)

    def terminate_child(self):
        try:
            if self.child_process:
                self.child_process.kill()
        except psutil.NoSuchProcess:
            logger.info('The process has already died.')
        return

    def isalive(self, process):
        """True if the child is still alive, false otherwise"""
        try:
            self.exitstatus = process.wait(timeout=0)
            return False
        except psutil.TimeoutExpired:
            return True

    def write(self, s):
        """Writes input into the child consoles input buffer."""

        if len(s) == 0:
            return 0
        if s[-1] == '\n':
            s = s[:-1]
        records = [self.createKeyEvent(c) for c in str(s)]
        if not self.consout:
            return ""

        # Store the current cursor position to hide characters in local echo disabled mode
        # (workaround).
        consinfo = self.consout.GetConsoleScreenBufferInfo()
        startCo = consinfo['CursorPosition']

        # Send the string to console input
        wrote = self.consin.WriteConsoleInput(records)

        # Wait until all input has been recorded by the console.
        ts = time.time()
        while self.consin.PeekConsoleInput(8) != ():
            if time.time() > ts + len(s) * .1 + .5:
                break
            time.sleep(.05)

        # Hide characters in local echo disabled mode (workaround).
        if not self.local_echo:
            self.consout.FillConsoleOutputCharacter(screenbufferfillchar, len(s), startCo)

        return wrote

    def createKeyEvent(self, char):
        """Creates a single key record corrosponding to
            the ascii character char."""

        evt = win32console.PyINPUT_RECORDType(win32console.KEY_EVENT)
        evt.KeyDown = True
        evt.Char = char
        evt.RepeatCount = 1
        return evt

    def initConsole(self, consout=None, window_size_x=80, window_size_y=25, buffer_size_x=80,
                    buffer_size_y=16000):
        if not consout:
            consout = self.getConsoleOut()

        self.consin = win32console.GetStdHandle(win32console.STD_INPUT_HANDLE)

        rect = win32console.PySMALL_RECTType(0, 0, window_size_x - 1, window_size_y - 1)
        consout.SetConsoleWindowInfo(True, rect)
        size = win32console.PyCOORDType(buffer_size_x, buffer_size_y)
        consout.SetConsoleScreenBufferSize(size)
        pos = win32console.PyCOORDType(0, 0)
        # Use NUL as fill char because it displays as whitespace
        # (if we interact() with the child)
        consout.FillConsoleOutputCharacter(screenbufferfillchar, size.X * size.Y, pos)

        consinfo = consout.GetConsoleScreenBufferInfo()
        self.__consSize = consinfo['Size']
        logger.info('self.__consSize: ' + str(self.__consSize))
        self.startCursorPos = consinfo['CursorPosition']

    def parseData(self, s):
        """Ensures that special characters are interpretted as
        newlines or blanks, depending on if there written over
        characters or screen-buffer-fill characters."""

        strlist = []
        for i, c in enumerate(s):
            if c == screenbufferfillchar:
                if (self.totalRead - self.lastRead + i + 1) % self.__consSize.X == 0:
                    strlist.append('\r\n')
            else:
                strlist.append(c)

        s = ''.join(strlist)
        return s

    def getConsoleOut(self):
        consfile = win32file.CreateFile(
            'CONOUT$',
            win32con.GENERIC_READ | win32con.GENERIC_WRITE,
            win32con.FILE_SHARE_READ | win32con.FILE_SHARE_WRITE,
            None,
            win32con.OPEN_EXISTING,
            0,
            0)

        self.consout = win32console.PyConsoleScreenBufferType(consfile)
        return self.consout

    def getCoord(self, offset):
        """Converts an offset to a point represented as a tuple."""

        x = offset % self.__consSize.X
        y = offset // self.__consSize.X
        return win32console.PyCOORDType(x, y)

    def getOffset(self, coord):
        """Converts a tuple-point to an offset."""

        return coord.X + coord.Y * self.__consSize.X

    def readConsole(self, startCo, endCo):
        """Reads the console area from startCo to endCo and returns it
        as a string."""

        if startCo is None:
            startCo = self.startCursorPos
            startCo.Y = startCo.Y

        if endCo is None:
            consinfo = self.consout.GetConsoleScreenBufferInfo()
            endCo = consinfo['CursorPosition']
            endCo = self.getCoord(0 + self.getOffset(endCo))

        buff = []
        self.lastRead = 0

        while True:
            startOff = self.getOffset(startCo)
            endOff = self.getOffset(endCo)
            readlen = endOff - startOff

            if readlen <= 0:
                break

            if readlen > 4000:
                readlen = 4000
            endPoint = self.getCoord(startOff + readlen)

            s = self.consout.ReadConsoleOutputCharacter(readlen, startCo)
            self.lastRead += len(s)
            self.totalRead += len(s)
            buff.append(s)

            startCo = endPoint

        return ''.join(buff)

    def readConsoleToCursor(self):
        """Reads from the current read position to the current cursor
        position and inserts the string into self.__buffer."""

        if not self.consout:
            return ""

        consinfo = self.consout.GetConsoleScreenBufferInfo()
        cursorPos = consinfo['CursorPosition']

        logger.spam('cursor: %r, current: %r' % (cursorPos, self.__currentReadCo))

        isSameX = cursorPos.X == self.__currentReadCo.X
        isSameY = cursorPos.Y == self.__currentReadCo.Y
        isSamePos = isSameX and isSameY

        logger.spam('isSameY: %r' % isSameY)
        logger.spam('isSamePos: %r' % isSamePos)

        if isSameY or not self.lastReadData.endswith('\r\n'):
            # Read the current slice again
            self.totalRead -= self.lastRead
            self.__currentReadCo.X = 0
            self.__currentReadCo.Y = self.__bufferY

        logger.spam('cursor: %r, current: %r' % (cursorPos, self.__currentReadCo))

        raw = self.readConsole(self.__currentReadCo, cursorPos)
        rawlist = []
        while raw:
            rawlist.append(raw[:self.__consSize.X])
            raw = raw[self.__consSize.X:]
        raw = ''.join(rawlist)
        s = self.parseData(raw)
        for i, line in enumerate(reversed(rawlist)):
            if line.endswith(screenbufferfillchar):
                # Record the Y offset where the most recent line break was detected
                self.__bufferY += len(rawlist) - i
                break

        logger.spam('lastReadData: %r' % self.lastReadData)
        if s:
            logger.debug('Read: %r' % s)
        else:
            logger.spam('Read: %r' % s)

        if isSamePos and self.lastReadData == s:
            logger.spam('isSamePos and self.lastReadData == s')
            s = ''

        if s:
            lastReadData = self.lastReadData
            pos = self.getOffset(self.__currentReadCo)
            self.lastReadData = s
            if isSameY or not lastReadData.endswith('\r\n'):
                # Detect changed lines
                self.__buffer.seek(pos)
                buf = self.__buffer.read()
                if raw.startswith(buf):
                    # Line has grown
                    rawslice = raw[len(buf):]
                    # Update last read bytes so line breaks can be detected in parseData
                    lastRead = self.lastRead
                    self.lastRead = len(rawslice)
                    s = self.parseData(rawslice)
                    self.lastRead = lastRead
                else:
                    # Cursor has been repositioned
                    s = '\r' + s
            self.__buffer.seek(pos)
            self.__buffer.truncate()
            self.__buffer.write(raw)

        self.__currentReadCo.X = cursorPos.X
        self.__currentReadCo.Y = cursorPos.Y

        return s

    def interact(self):
        """Displays the child console for interaction."""

        logger.debug('Start interact window')
        win32gui.ShowWindow(win32console.GetConsoleWindow(), win32con.SW_SHOW)

    def sendeof(self):
        """This sends an EOF to the host. This sends a character which inform the host that child
        has been finished, and all of it's output has been send to host.
        """

        self.send_to_host(EOF_CHAR)


class ConsoleReaderSocket(ConsoleReaderBase):

    def create_connection(self, **kwargs):
        try:
            self.port = kwargs['port']
            # Create a TCP/IP socket
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            server_address = ('localhost', self.port)
            self.sock.bind(server_address)
            logger.info(f'Socket started at port: {self.port}')

            # Listen for incoming connections
            self.sock.settimeout(5)
            self.sock.listen(1)
            self.connection, client_address = self.sock.accept()
            self.connection.settimeout(.01)
            logger.info(f'Client connected: {client_address}')
        except Exception as e:  # pragma: no cover
            # I hope this code is unreachable.
            logger.error(f"Port: {self.port} {e}")
            raise

    def close_connection(self):
        if self.connection:
            self.connection.shutdown(socket.SHUT_RDWR)
            self.connection.close()
            self.connection = None

    def send_to_host(self, msg):
        # convert to bytes
        if isinstance(msg, str):
            msg = str.encode(msg)
        if msg:
            logger.debug(f'Sending msg: {msg}')
        else:
            logger.spam(f'Sending msg: {msg}')
        self.connection.sendall(msg)

    def get_from_host(self):
        try:
            msg = self.connection.recv(4096)
        except socket.timeout as e:
            err = e.args[0]
            # this next if/else is a bit redundant, but illustrates how the
            # timeout exception is setup
            if err == 'timed out':
                logger.debug('recv timed out, retry later')
                return b''
            else:
                raise
        else:
            if len(msg) == 0:
                raise Exception('orderly shutdown on server end')
            else:
                # got a message do something :)
                return msg


class ConsoleReaderPipe(ConsoleReaderBase):
    def create_connection(self, timeout=-1, **kwargs):
        if timeout == -1:
            timeout = self.timeout
        if timeout is None:
            end_time = float('inf')
        else:
            end_time = time.time() + timeout

        pipe_name = 'wexpect_{}'.format(self.console_pid)
        pipe_full_path = r'\\.\pipe\{}'.format(pipe_name)
        logger.info('Start pipe server: %s', pipe_full_path)
        self.pipe = win32pipe.CreateNamedPipe(
            pipe_full_path,
            win32pipe.PIPE_ACCESS_DUPLEX,
            win32pipe.PIPE_TYPE_MESSAGE | win32pipe.PIPE_READMODE_MESSAGE | win32pipe.PIPE_NOWAIT,
            1, 65536, 65536, 10000, None)
        logger.info("waiting for client")
        while True:
            if end_time < time.time():
                raise TIMEOUT('Connect to child has been timed out.')
            try:
                win32pipe.ConnectNamedPipe(self.pipe, None)
                break
            except Exception as e:
                logger.debug(e)
                time.sleep(0.2)
        logger.info('got client')

    def close_connection(self):
        if self.pipe:
            win32file.CloseHandle(self.pipe)

    def send_to_host(self, msg):
        # convert to bytes
        if isinstance(msg, str):
            msg = str.encode(msg)
        if msg:
            logger.debug(f'Sending msg: {msg}')
        else:
            logger.spam(f'Sending msg: {msg}')
        win32file.WriteFile(self.pipe, msg)

    def get_from_host(self):
        data, avail, bytes_left = win32pipe.PeekNamedPipe(self.pipe, 4096)
        logger.spam(f'data: {data}  avail:{avail}  bytes_left{bytes_left}')
        if avail > 0:
            resp = win32file.ReadFile(self.pipe, 4096)
            ret = resp[1]
            return ret
        else:
            return b''