contrib/automation/hgautomation/winrm.py
changeset 42024 b05a3e28cf24
child 42279 f30184484dd1
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/contrib/automation/hgautomation/winrm.py	Fri Mar 15 11:24:08 2019 -0700
@@ -0,0 +1,82 @@
+# winrm.py - Interact with Windows Remote Management (WinRM)
+#
+# Copyright 2019 Gregory Szorc <gregory.szorc@gmail.com>
+#
+# 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=120, 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_cmd('echo "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 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)))