AwaitDriver.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

using System.Diagnostics;
using System.Runtime.InteropServices;

namespace AwaitDriver
{
    public class AwaitDriver
    {
        /// <summary>
        /// Creates a new instance of the AwaitDriver class, which lets you perform
        /// expect-like functionality on Console applications.
        /// </summary>
        public AwaitDriver()
        {
            // Launch a new instance of PowerShell and steal its console input / output handles
            ProcessStartInfo startInfo = new ProcessStartInfo("C:\\windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", "-NoProfile");
            startInfo.WindowStyle = ProcessWindowStyle.Hidden;
            driverHost = Process.Start(startInfo);

            System.Threading.Thread.Sleep(100);

            NativeMethods.FreeConsole();
            bool result = NativeMethods.AttachConsole(driverHost.Id);

            const UInt32 STD_INPUT_HANDLE = 0xFFFFFFF6;
            const UInt32 STD_OUTPUT_HANDLE = 0xFFFFFFF5;
            hInput = NativeMethods.GetStdHandle(STD_INPUT_HANDLE);
            hOutput = NativeMethods.GetStdHandle(STD_OUTPUT_HANDLE);
        }

        // Handles from the driver process
        Process driverHost = null;
        IntPtr hInput = IntPtr.Zero;
        IntPtr hOutput = IntPtr.Zero;

        // State to remember information between Read() calls so
        // that we don't have to scan the full buffer each time.
        int lastReadPosition = 0;
        List<string> lastRawContent = new List<string>();

        /// <summary>
        /// 
        /// Sends the specified input to the driver process. Does not
        /// include a newline at the end of the input.
        /// 
        /// Input can be basic characters, or can use the same set of metacharacters
        /// that the SendKeys() API supports, such as {ESC}.
        /// 
        /// http://msdn.microsoft.com/en-us/library/system.windows.forms.sendkeys.send(v=vs.110).aspx
        /// 
        /// While the syntax is the same as the SendKeys() API, the approach is not
        /// based on SendKeys() at all - to ensure that applications can be automated while
        /// you use your keyboard and mouse for other things.
        /// 
        /// </summary>
        /// <param name="text">The text to send</param>
        public void Send(string text)
        {
            List<NativeMethods.INPUT_RECORD> inputs = new List<NativeMethods.INPUT_RECORD>();

            foreach (string inputElement in SendKeysParser.Parse(text))
            {
                foreach (NativeMethods.INPUT_RECORD mappedInput in SendKeysParser.MapInput(inputElement))
                {
                    inputs.Add(mappedInput);
                }
            }

            uint eventsWritten = 0;
            NativeMethods.WriteConsoleInput(hInput, inputs.ToArray(), (uint)inputs.Count, out eventsWritten);
        }

        /// <summary>
        /// Sends the specified input to the driver process, including a trailing newline.
        /// </summary>
        /// <param name="text">The text to send</param>
        public void SendLine(string input)
        {
            Send(input + "{ENTER}");
        }

        /// <summary>
        /// Waits for the given string to be available as a command response, and
        /// returns the result of that reponse.
        /// 
        /// This method is stateful. Once output has been returned from an AwaitOutput()
        /// call, it will not be output again. Further AwaitOutput() calls will operate
        /// against output produced after preceeding AwaitOutput() calls complete.
        /// 
        /// </summary>
        /// <param name="expected">The string to search for in the command response.</param>
        public string AwaitOutput(string expected)
        {
            return AwaitOutput(expected, false);
        }

        /// <summary>
        /// Waits for the given string to be available as a command response, and
        /// returns the result of that reponse.
        /// 
        /// </summary>
        /// <param name="expected">The string to search for in the command response.</param>
        /// <param name="all">
        /// If 'All' is specified, this method acts against the entire content of the
        /// console buffer, and the output of this method is the entire content of the
        /// console buffer.
        /// </param>
        public string AwaitOutput(string expected, bool all)
        {
            // Save the internal buffer position state so that we can restore
            // it if the content doesn't match.
            int savedStartPosition = lastReadPosition;
            List<string> savedLastRawContent = new List<string>(lastRawContent);
            string output = null;

            // Should be able to replace this with a WinEventHook for console changes:
            // http://msdn.microsoft.com/en-us/library/windows/desktop/ms682102(v=vs.85).aspx
            while (true)
            {
                output = ReadOutput(all);

                if (output.Contains(expected))
                {
                    break;
                }
                else
                {
                    lastReadPosition = savedStartPosition;
                    lastRawContent = new List<string>(savedLastRawContent);
                }

                System.Threading.Thread.Sleep(10);
            }

            return output;
        }

        /// <summary>
        /// Closes the AwaitDriver class, which lets you perform
        /// expect-like functionality on Console applications.
        /// </summary>
        public void Close()
        {
            driverHost.Kill();
        }

        public string ReadOutput()
        {
            return ReadOutput(false);
        }

        // Can scrape 650 lines of output at 11ms per scrape.
        public string ReadOutput(bool all)
        {
            // Check the current console screen buffer dimensions, cursor coordinates / etc.
            NativeMethods.CONSOLE_SCREEN_BUFFER_INFO csbi;
            NativeMethods.GetConsoleScreenBufferInfo(hOutput, out csbi);

            StringBuilder output = new StringBuilder();

            // If the cursor has gone before where we last scanned, then the screen
            // has been cleared and we should reset our state.
            if (lastReadPosition > csbi.dwCursorPosition.Y + 1)
            {
                ResetLastScanInfo();
            }

            // If the cursor is at the end of the buffer, then we no longer have a
            // quick way to determine if new content has been written. To address this,
            // every ReadOutput() call retains the last 10 lines of output that it read.
            //
            // As we scan through the buffer, we fill the 'currentHeuristicContent' list
            // with the output lines we see. Once we have 10 matching lines of output, we
            // know we're at the end of the previous scan.
            //
            // Additionally, we only do this heuristic scan if they haven't specified the
            // 'All' flag.
            bool inHeuristicContentScan = false;
            List<string> currentHeuristicContent = new List<string>();

            if ((!all) && (lastReadPosition == csbi.dwSize.Y))
            {
                inHeuristicContentScan = true;
                lastReadPosition = 0;
            }

            // Figure out where to start and stop scanning
            int startReadPosition = lastReadPosition;
            int endReadPosition = csbi.dwCursorPosition.Y;

            if (all)
            {
                startReadPosition = 0;
                endReadPosition = csbi.dwSize.Y - 1;
            }

            // Go through each line in the buffer
            for (int row = startReadPosition; row <= endReadPosition; row++)
            {
                int width = csbi.dwMaximumWindowSize.X;
                StringBuilder lpCharacter = new StringBuilder(width - 1);

                // Read the current line from the buffer
                NativeMethods.COORD dwReadCoord;
                dwReadCoord.X = 0;
                dwReadCoord.Y = (short)row;

                uint lpNumberOfCharsRead = 0;
                NativeMethods.ReadConsoleOutputCharacter(hOutput, lpCharacter, (uint)width, dwReadCoord, out lpNumberOfCharsRead);

                // If we're in a heuristic scan, and the current line to the heuristic scan buffer
                if (inHeuristicContentScan)
                {
                    currentHeuristicContent.Add(lpCharacter.ToString());

                    // Ensure the two lists remain the same size
                    if (currentHeuristicContent.Count > lastRawContent.Count)
                    {
                        currentHeuristicContent.RemoveAt(0);
                    }

                    // If the heuristic scan buffer is the same size as the trailing buffer from the last ReadOutput()
                    // call, we can see if they match.
                    if (currentHeuristicContent.Count == lastRawContent.Count)
                    {
                        bool foundContentMatch = true;

                        // Only scan to the second-last line of the saved content buffer, as the last line
                        // is frequently different. For example, a prompt in the last scan, and now a prompt
                        // that includes the user's command.
                        for (int contentIndex = 0; contentIndex < currentHeuristicContent.Count - 1; contentIndex++)
                        {
                            if (lastRawContent[contentIndex] != currentHeuristicContent[contentIndex])
                            {
                                foundContentMatch = false;
                                break;
                            }
                        }

                        // All 10 lines matched
                        if (foundContentMatch)
                        {
                            inHeuristicContentScan = false;
                            continue;
                        }
                    }
                }

                // If we're not just scanning through the buffer heuristically, then capture the content.
                if (!inHeuristicContentScan)
                {
                    lastRawContent.Add(lpCharacter.ToString());

                    if (lastRawContent.Count > 10)
                    {
                        lastRawContent.RemoveAt(0);
                    }

                    output.AppendLine(lpCharacter.ToString().Substring(0, width).TrimEnd());
                }
            }

            // Update our state to remember where to start scanning the buffer next.
            lastReadPosition = endReadPosition + 1;

            // If we were at the end of the buffer and made it all the way here without matching,
            // then our heuristics (i.e.: detecting clearing the screen) got broken.
            // This can happen with repetitive commands when the buffer is full, such as:

            // Get-Date
            // cls
            // Get-Date

            // Reset them.
            if (inHeuristicContentScan)
            {
                ResetLastScanInfo();
                return ReadOutput(all);
            }
            else
            {
                // We got content - return it.
                return output.ToString().TrimEnd();
            }
        }

        private void ResetLastScanInfo()
        {
            lastReadPosition = 0;
            lastRawContent.Clear();
        }
    }

    class SendKeysParser
    {
        public static List<string> Parse(string input)
        {
            List<string> output = new List<string>();

            bool scanningKeyName = false;
            StringBuilder keyNameBuffer = new StringBuilder();

            // Iterate through the string
            for (int index = 0; index < input.Length; index++)
            {
                // Save the current item
                char currentChar = input[index];

                // We may have the start of a command
                if (currentChar == '{')
                {
                    if (scanningKeyName)
                    {
                        throw new Exception("The character '{' is not a valid in a key name. " +
                            "To include the '{' character in your text, escape it with another: {{.");
                    }

                    // If it's escaped, then add it to output.
                    if ((index < (input.Length - 1)) &&
                        (input[index + 1] == '{'))
                    {
                        output.Add(currentChar.ToString());
                        index++;
                    }
                    else
                    {
                        // Otherwise, we found the start of a key name.
                        scanningKeyName = true;
                    }
                }
                else if (currentChar == '}')
                {
                    // We may have the end of a key name

                    // If it's escaped, then add it to output.
                    if ((index < (input.Length - 1)) &&
                        (input[index + 1] == '}'))
                    {
                        // But not if we're scanning a key name
                        if (scanningKeyName)
                        {
                            throw new Exception("The character '}' is not a valid in a key name. " +
                                "To include the '}' character in your text, escape it with another: }}.");
                        }

                        output.Add(currentChar.ToString());
                        index++;
                    }
                    else
                    {
                        // Not escaped

                        // If we're scanning a key name, record it.
                        if (scanningKeyName)
                        {
                            string keyName = keyNameBuffer.ToString();
                            if (String.IsNullOrEmpty(keyName))
                            {
                                throw new Exception("Key names may not be empty.");
                            }

                            output.Add(keyNameBuffer.ToString());
                            scanningKeyName = false;
                        }
                        else
                        {
                            throw new Exception("The character '}' is not a valid by itself. " +
                                "To include the '}' character in your text, escape it with another: }}.");
                        }
                    }
                }
                else
                {
                    // Just a letter
                    if (scanningKeyName)
                    {
                        keyNameBuffer.Append(currentChar);
                    }
                    else
                    {
                        output.Add(currentChar.ToString());
                    }
                }
            }

            // We got to the end of the string.
            if (scanningKeyName)
            {
                throw new Exception("The character '{' (representing the start of a key name) did not have a matching '}' " +
                    "character. To include the '{' character in your text, escape it with another: {{.");
            }

            return output;
        }

        internal static List<NativeMethods.INPUT_RECORD> MapInput(string inputElement)
        {
            List<NativeMethods.INPUT_RECORD> inputs = new List<NativeMethods.INPUT_RECORD>();

            NativeMethods.INPUT_RECORD input = new NativeMethods.INPUT_RECORD();
            input.EventType = 0x0001;

            NativeMethods.KEY_EVENT_RECORD keypress = new NativeMethods.KEY_EVENT_RECORD();
            keypress.dwControlKeyState = 0;
            keypress.wRepeatCount = 1;

            // Just a regular character
            if (inputElement.Length == 1)
            {
                keypress.UnicodeChar = inputElement[0];
            }
            else
            {
                switch (inputElement.ToUpperInvariant())
                {
                    case "BACKSPACE":
                    case "BS":
                    case "BKSP":
                        keypress = GetKeyPressForSimpleKey(keypress, 0x08);
                        break;

                    case "BREAK":
                        keypress = GetKeyPressForSimpleKey(keypress, 0x03);
                        keypress.dwControlKeyState = (uint) NativeMethods.ControlKeyStates.LEFT_CTRL_PRESSED;
                        keypress.UnicodeChar = (char) 0;
                        break;

                    case "ENTER":
                        keypress = GetKeyPressForSimpleKey(keypress, 0x0D);
                        break;

                    case "ESC":
                        keypress = GetKeyPressForSimpleKey(keypress, 0x1B);
                        break;
                }
            }

            keypress.bKeyDown = true;
            input.KeyEvent = keypress;
            inputs.Add(input);

            keypress.bKeyDown = false;
            keypress.dwControlKeyState = 0;
            input.KeyEvent = keypress;
            inputs.Add(input);

            return inputs;
        }

        private static NativeMethods.KEY_EVENT_RECORD GetKeyPressForSimpleKey(NativeMethods.KEY_EVENT_RECORD keypress, uint uCode)
        {
            keypress.UnicodeChar = (char)uCode;
            keypress.wVirtualKeyCode = (ushort)uCode;
            keypress.wVirtualScanCode = (ushort)NativeMethods.MapVirtualKey(uCode, MAPVK_VK_TO_VSC);
            return keypress;
        }

        const ushort MAPVK_VK_TO_VSC = 0x00;
    }

    class NativeMethods
    {
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool AttachConsole(int pid);

        [DllImport("kernel32.dll")]
        internal static extern bool FreeConsole();

        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool WriteConsoleInput(IntPtr hConsoleInput, INPUT_RECORD[] lpBuffer, uint nLength, out uint lpNumberOfEventsWritten);

        [DllImport("user32.dll")]
        internal static extern uint MapVirtualKey(uint uCode, uint uMapType);

        [DllImport("kernel32.dll")]
        internal static extern IntPtr GetStdHandle(uint nStdHandle);

        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool FlushConsoleInputBuffer(IntPtr hConsoleInput);

        [DllImport("Kernel32")]
        internal static extern bool ReadConsoleOutputCharacter(IntPtr hConsoleOutput, StringBuilder lpCharacter, uint nLength, COORD dwReadCoord, out uint lpNumberOfCharsRead);

        [DllImport("kernel32.dll")]
        internal static extern bool GetConsoleScreenBufferInfo(IntPtr hConsoleOutput, out CONSOLE_SCREEN_BUFFER_INFO lpConsoleScreenBufferInfo);

        [StructLayout(LayoutKind.Sequential)]
        internal struct COORD
        {
            internal short X;
            internal short Y;
        }

        internal struct SMALL_RECT
        {
            internal short Left;
            internal short Top;
            internal short Right;
            internal short Bottom;
        }

        internal struct CONSOLE_SCREEN_BUFFER_INFO
        {
            internal COORD dwSize;
            internal COORD dwCursorPosition;
            internal short wAttributes;
            internal SMALL_RECT srWindow;
            internal COORD dwMaximumWindowSize;
        }

        [StructLayout(LayoutKind.Explicit)]
        internal struct INPUT_RECORD
        {
            [FieldOffset(0)]
            internal ushort EventType;
            [FieldOffset(4)]
            internal KEY_EVENT_RECORD KeyEvent;
        };

        [StructLayout(LayoutKind.Explicit, CharSet = CharSet.Unicode)]
        internal struct KEY_EVENT_RECORD
        {
            [FieldOffset(0), MarshalAs(UnmanagedType.Bool)]
            internal bool bKeyDown;
            [FieldOffset(4), MarshalAs(UnmanagedType.U2)]
            internal ushort wRepeatCount;
            [FieldOffset(6), MarshalAs(UnmanagedType.U2)]
            internal ushort wVirtualKeyCode;
            [FieldOffset(8), MarshalAs(UnmanagedType.U2)]
            internal ushort wVirtualScanCode;
            [FieldOffset(10)]
            internal char UnicodeChar;
            [FieldOffset(12), MarshalAs(UnmanagedType.U4)]
            internal uint dwControlKeyState;
        }

        internal enum ControlKeyStates
        {
            RIGHT_ALT_PRESSED = 0x1,
            LEFT_ALT_PRESSED = 0x2,
            RIGHT_CTRL_PRESSED = 0x4,
            LEFT_CTRL_PRESSED = 0x8,
            SHIFT_PRESSED = 0x10,
            NUMLOCK_ON = 0x20,
            SCROLLLOCK_ON = 0x40,
            CAPSLOCK_ON = 0x80,
            ENHANCED_KEY = 0x100
        }
    }
}