contrib/automation/hgautomation/winrm.py
author Gregory Szorc <gregory.szorc@gmail.com>
Fri, 19 Apr 2019 06:07:00 -0700
changeset 42279 f30184484dd1
parent 42024 b05a3e28cf24
child 42661 e91930d712e8
permissions -rw-r--r--
automation: wait longer for WinRM connection I got a few timeouts waiting for only 120s for the WinRM connection to become available. Increasing to 180s seems to fix. I guess AWS isn't as consistent as I would like :( Differential Revision: https://phab.mercurial-scm.org/D6287
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
42024
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     1
# winrm.py - Interact with Windows Remote Management (WinRM)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     2
#
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     3
# Copyright 2019 Gregory Szorc <gregory.szorc@gmail.com>
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     4
#
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     5
# This software may be used and distributed according to the terms of the
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     6
# GNU General Public License version 2 or any later version.
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     7
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     8
# no-check-code because Python 3 native.
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
     9
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    10
import logging
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    11
import pprint
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    12
import time
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    13
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    14
from pypsrp.client import (
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    15
    Client,
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    16
)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    17
from pypsrp.powershell import (
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    18
    PowerShell,
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    19
    PSInvocationState,
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    20
    RunspacePool,
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    21
)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    22
import requests.exceptions
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    23
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    24
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    25
logger = logging.getLogger(__name__)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    26
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    27
42279
f30184484dd1 automation: wait longer for WinRM connection
Gregory Szorc <gregory.szorc@gmail.com>
parents: 42024
diff changeset
    28
def wait_for_winrm(host, username, password, timeout=180, ssl=False):
42024
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    29
    """Wait for the Windows Remoting (WinRM) service to become available.
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    30
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    31
    Returns a ``psrpclient.Client`` instance.
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    32
    """
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    33
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    34
    end_time = time.time() + timeout
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    35
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    36
    while True:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    37
        try:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    38
            client = Client(host, username=username, password=password,
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    39
                            ssl=ssl, connection_timeout=5)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    40
            client.execute_cmd('echo "hello world"')
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    41
            return client
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    42
        except requests.exceptions.ConnectionError:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    43
            if time.time() >= end_time:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    44
                raise
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    45
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    46
            time.sleep(1)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    47
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    48
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    49
def format_object(o):
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    50
    if isinstance(o, str):
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    51
        return o
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    52
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    53
    try:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    54
        o = str(o)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    55
    except TypeError:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    56
        o = pprint.pformat(o.extended_properties)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    57
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    58
    return o
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    59
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    60
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    61
def run_powershell(client, script):
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    62
    with RunspacePool(client.wsman) as pool:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    63
        ps = PowerShell(pool)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    64
        ps.add_script(script)
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    65
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    66
        ps.begin_invoke()
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    67
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    68
        while ps.state == PSInvocationState.RUNNING:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    69
            ps.poll_invoke()
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    70
            for o in ps.output:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    71
                print(format_object(o))
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    72
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    73
            ps.output[:] = []
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    74
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    75
        ps.end_invoke()
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    76
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    77
        for o in ps.output:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    78
            print(format_object(o))
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    79
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    80
        if ps.state == PSInvocationState.FAILED:
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    81
            raise Exception('PowerShell execution failed: %s' %
b05a3e28cf24 automation: perform tasks on remote machines
Gregory Szorc <gregory.szorc@gmail.com>
parents:
diff changeset
    82
                            ' '.join(map(format_object, ps.streams.error)))