view contrib/automation/hgautomation/winrm.py @ 46513:c3c7a86e9c24

tests: fix differing output between py2 and py3

db9e33beb0fb broke the tests because of the difference in bytestring repr
between py2 and py3. Rather than backout that change for so little, I figured
I'd fix it myself.
Hopefully Python 2 supports gets dropped very soon.

Differential Revision: https://phab.mercurial-scm.org/D9987

# winrm.py - Interact with Windows Remote Management (WinRM)
#
# Copyright 2019 Gregory Szorc <[e-mail]>
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.

# no-check-code because Python 3 native.

import logging
import pprint
import time

from pypsrp.client import Client
from pypsrp.powershell import (
    PowerShell,
    PSInvocationState,
    RunspacePool,
)
import requests.exceptions


logger = logging.getLogger(__name__)


def wait_for_winrm(host, username, password, timeout=180, ssl=False):
    """Wait for the Windows Remoting (WinRM) service to become available.

    Returns a ``psrpclient.Client`` instance.
    """

    end_time = time.time() + timeout

    while True:
        try:
            client = Client(
                host,
                username=username,
                password=password,
                ssl=ssl,
                connection_timeout=5,
            )
            client.execute_ps("Write-Host 'Hello, World!'")
            return client
        except requests.exceptions.ConnectionError:
            if time.time() >= end_time:
                raise

            time.sleep(1)


def format_object(o):
    if isinstance(o, str):
        return o

    try:
        o = str(o)
    except (AttributeError, TypeError):
        o = pprint.pformat(o.extended_properties)

    return o


def run_powershell(client, script):
    with RunspacePool(client.wsman) as pool:
        ps = PowerShell(pool)
        ps.add_script(script)

        ps.begin_invoke()

        while ps.state == PSInvocationState.RUNNING:
            ps.poll_invoke()
            for o in ps.output:
                print(format_object(o))

            ps.output[:] = []

        ps.end_invoke()

        for o in ps.output:
            print(format_object(o))

        if ps.state == PSInvocationState.FAILED:
            raise Exception(
                'PowerShell execution failed: %s'
                % ' '.join(map(format_object, ps.streams.error))
            )