add a retry option when the connection fails

contents/winrm-exec.py

@@ -135,6 +135,9 @@ def httpclient_log(*args):
 exitBehaviour = "console"
 cleanescapingflg = False
 enabledHttpDebug = False
+retryconnection = 1
+retryconnectiondelay = 0
+username = None
 
 if "RD_CONFIG_AUTHTYPE" in os.environ:
     authentication = os.getenv("RD_CONFIG_AUTHTYPE")
@@ -187,6 +190,12 @@ def httpclient_log(*args):
     else:
         enabledHttpDebug = False
 
+if "RD_CONFIG_RETRYCONNECTION" in os.environ:
+    retryconnection = int(os.getenv("RD_CONFIG_RETRYCONNECTION"))
+
+if "RD_CONFIG_RETRYCONNECTIONDELAY" in os.environ:
+    retryconnectiondelay = int(os.getenv("RD_CONFIG_RETRYCONNECTIONDELAY"))
+
 exec_command = os.getenv("RD_EXEC_COMMAND")
 log.debug("Command will be executed: " + exec_command)
 
@@ -243,8 +252,8 @@ def httpclient_log(*args):
 log.debug("username:" + username)
 log.debug("nossl:" + str(nossl))
 log.debug("diabletls12:" + str(diabletls12))
-log.debug("krb5config:" + krb5config)
-log.debug("kinit command:" + kinit)
+log.debug("krb5config:" + str(krb5config))
+log.debug("kinit command:" + str(kinit))
 log.debug("kerberos delegation:" + str(krbdelegation))
 log.debug("shell:" + shell)
 log.debug("output_charset:" + output_charset)
@@ -253,6 +262,8 @@ def httpclient_log(*args):
 log.debug("exit Behaviour:" + exitBehaviour)
 log.debug("cleanescapingflg: " + str(cleanescapingflg))
 log.debug("enabledHttpDebug: " + str(enabledHttpDebug))
+log.debug("retryConnection: " + str(retryconnection))
+log.debug("retryConnectionDelay: " + str(retryconnectiondelay))
 log.debug("------------------------------------------")
 
 if enabledHttpDebug:
@@ -314,7 +325,7 @@ def httpclient_log(*args):
 winrm.Session._clean_error_msg = winrm_session._clean_error_msg
 winrm.Session._strip_namespace = winrm_session._strip_namespace
 
-tsk = winrm_session.RunCommand(session, shell, exec_command, output_charset)
+tsk = winrm_session.RunCommand(session, shell, exec_command, retryconnection, retryconnectiondelay, output_charset)
 t = threading.Thread(target=tsk.get_response)
 t.start()
 realstdout = sys.stdout

contents/winrm-filecopier.py

@@ -16,8 +16,9 @@
 from colored_formatter import ColoredFormatter
 import kerberosauth
 import http.client
+import winrm_session
 
-#checking and importing dependencies
+# checking and importing dependencies
 ISPY3 = sys.version_info[0] == 3
 WINRM_INSTALLED = False
 URLLIB_INSTALLED = False
@@ -170,8 +171,10 @@ class WinRmError(RemoteCommandError):
 
 class CopyFiles(object):
 
-    def __init__(self, session):
-        self.session=session
+    def __init__(self, session, retry, retry_delay):
+        self.session = session
+        self.retry = retry
+        self.retry_delay = retry_delay
 
 
     def winrm_upload(self,
@@ -189,10 +192,12 @@ def winrm_upload(self,
 
         print("coping file %s to %s" % (local_path, full_path))
 
-        self.session.run_ps('if (!(Test-Path {0})) {{ New-Item -ItemType directory -Path {0} }}'.format(remote_path))
+        self.session.run_ps('if (!(Test-Path {0})) {{ New-Item -ItemType directory -Path {0} }}'.format(remote_path),
+                            retry=self.retry,
+                            retry_delay=self.retry_delay)
 
-        if(override):
-            self.session.run_ps('if ((Test-Path {0} -PathType Leaf)) {{ rm {0} }}'.format(full_path))
+        if override:
+            self.session.run_ps('if ((Test-Path {0} -PathType Leaf)) {{ rm {0} }}'.format(full_path), retry=self.retry, retry_delay=self.retry_delay)
 
         size = os.stat(local_path).st_size
         with open(local_path, 'rb') as f:
@@ -258,6 +263,10 @@ def winrm_upload(self,
 enabledHttpDebug = False
 readtimeout = None
 operationtimeout = None
+retryconnection = 1
+retryconnectiondelay = 0
+certpath = None
+username = None
 
 if os.environ.get('RD_CONFIG_OVERRIDE') == 'true':
     override = True
@@ -333,12 +342,17 @@ def winrm_upload(self,
     else:
         enabledHttpDebug = False
 
+if "RD_CONFIG_RETRYCONNECTION" in os.environ:
+    retryconnection = int(os.getenv("RD_CONFIG_RETRYCONNECTION"))
+
+if "RD_CONFIG_RETRYCONNECTIONDELAY" in os.environ:
+    retryconnectiondelay = int(os.getenv("RD_CONFIG_RETRYCONNECTIONDELAY"))
+
 if enabledHttpDebug:
     httpclient_logging_patch(logging.DEBUG)
 
 endpoint = transport+'://'+args.hostname+':'+port
-arguments = {}
-arguments["transport"] = authentication
+arguments = {"transport": authentication}
 
 if(nossl == True):
     arguments["server_cert_validation"] = "ignore"
@@ -388,10 +402,12 @@ def winrm_upload(self,
                         auth=(username, password),
                         **arguments)
 
+winrm.Session.run_cmd = winrm_session.run_cmd
+winrm.Session.run_ps = winrm_session.run_ps
+winrm.Session._clean_error_msg = winrm_session._clean_error_msg
+winrm.Session._strip_namespace = winrm_session._strip_namespace
 
-winrm.Session._clean_error_msg = _clean_error_msg
-
-copy = CopyFiles(session)
+copy = CopyFiles(session, retryconnection, retryconnectiondelay)
 
 destination = args.destination
 filename = ntpath.basename(args.destination)

contents/winrm_session.py

@@ -1,10 +1,14 @@
 from __future__ import unicode_literals
 import xml.etree.ElementTree as ET
+
 try:
-	import os; os.environ['PATH']
+    import os;
+
+    os.environ['PATH']
 except:
-	import os
-	os.environ.setdefault('PATH', '')
+    import os
+
+    os.environ.setdefault('PATH', '')
 try:
     from StringIO import StringIO
 except ImportError:
@@ -24,6 +28,9 @@
 import sys
 import types
 import re
+import time
+import logging
+from requests import ConnectionError
 
 PY2 = sys.version_info[0] == 2
 PY3 = sys.version_info[0] == 3
@@ -46,20 +53,40 @@
     text_type = unicode
     binary_type = str
 
+log = logging.getLogger()
 
 # TODO: this PR https://github.com/diyan/pywinrm/pull/55 will add this fix.
 # when this PR is merged, this won't be needed anymore
-def run_cmd(self, command, args=(), out_stream=None, err_stream=None, output_charset='utf-8'):
+def run_cmd(self, command, args=(), out_stream=None, err_stream=None, retry=1, retry_delay=0, output_charset='utf-8'):
     self.protocol.get_command_output = protocol.get_command_output
     winrm.Session._clean_error_msg = self._clean_error_msg
 
+    retryCount = 0
+
     envs = {}
     for a in os.environ:
         if a.startswith(RD) and INVALID_CHAR not in os.environ[a]:
-           envs.update({a:os.getenv(a)})
+            envs.update({a: os.getenv(a)})
+
+    shell_id = None
+
+    while retryCount < retry:
+        try:
+            shell_id = self.protocol.open_shell(codepage=65001, env_vars=envs)
+            break
+        except ConnectionError as e:
+            if retryCount < retry:
+                retryCount += 1
+                log.debug("error connecting " + str(e))
+                log.debug("Retrying connection " + str(retryCount) + "/" + str(retry) + " in " + str(retry_delay) + " seconds")
+                time.sleep(retry_delay)
+            else:
+                break
+
+
+    if shell_id is None:
+        raise Exception("Connection failed after {0} retries".format(retry))
 
-    # TODO optimize perf. Do not call open/close shell every time
-    shell_id = self.protocol.open_shell(codepage=65001, env_vars=envs)
     command_id = self.protocol.run_command(shell_id, command, args)
     rs = Response(self.protocol.get_command_output(self.protocol, shell_id, command_id, out_stream, err_stream))
 
@@ -71,19 +98,24 @@ def run_cmd(self, command, args=(), out_stream=None, err_stream=None, output_cha
     return rs
 
 
-def run_ps(self, script, out_stream=None, err_stream=None, output_charset='utf-8'):
+def run_ps(self, script, out_stream=None, err_stream=None, retry=1, retry_delay=0, output_charset='utf-8'):
     """base64 encodes a Powershell script and executes the powershell
     encoded script command
     """
     script = to_text(script)
     encoded_ps = base64.b64encode(script.encode('utf_16_le')).decode('ascii')
-    rs = self.run_cmd('powershell -encodedcommand {0}'.format(encoded_ps),out_stream=out_stream, err_stream=err_stream, output_charset=output_charset)
+    rs = self.run_cmd('powershell -encodedcommand {0}'.format(encoded_ps),
+                      out_stream=out_stream,
+                      err_stream=err_stream,
+                      output_charset=output_charset,
+                      retry=retry,
+                      retry_delay=retry_delay)
 
     return rs
 
 
 def _clean_error_msg(self, msg, encoding='utf-8'):
-    #data=""
+    # data=""
     msg = to_text(msg, encoding)
 
     if msg.startswith("#< CLIXML") or "<Objs Version=" in msg or "-1</PI><PC>" in msg:
@@ -114,6 +146,7 @@ def _clean_error_msg(self, msg, encoding='utf-8'):
 
     return msg
 
+
 def _strip_namespace(self, xml):
     """strips any namespaces from an xml string"""
     value = to_bytes(xml)
@@ -123,8 +156,10 @@ def _strip_namespace(self, xml):
         value = value.replace(match.group(), b"")
     return value
 
+
 class Response(object):
     """Response from a remote command execution"""
+
     def __init__(self, args):
         self.std_out, self.std_err, self.status_code = args
 
@@ -135,32 +170,36 @@ def __repr__(self):
 
 
 class RunCommand:
-    def __init__(self, session, shell, command, output_charset='utf-8'):
+    def __init__(self, session, shell, command, retry, retry_delay, output_charset='utf-8'):
         self.stat, self.o_std, self.e_std = None, None, None
         self.o_stream = BytesIO()
         self.e_stream = BytesIO()
         self.session = session
         self.exec_command = command
         self.shell = shell
         self.output_charset = output_charset
+        self.retry = retry
+        self.retry_delay = retry_delay
 
     def get_response(self):
         try:
             if self.shell == "cmd":
-                response = self.session.run_cmd(self.exec_command, out_stream=self.o_stream, err_stream=self.e_stream, output_charset=self.output_charset)
+                response = self.session.run_cmd(self.exec_command, out_stream=self.o_stream, err_stream=self.e_stream,
+                                                output_charset=self.output_charset, retry=self.retry, retry_delay=self.retry_delay)
                 self.o_std = response.std_out
                 self.e_std = response.std_err
                 self.stat = response.status_code
 
             if self.shell == "powershell":
-                response = self.session.run_ps(self.exec_command, out_stream=self.o_stream, err_stream=self.e_stream, output_charset=self.output_charset)
+                response = self.session.run_ps(self.exec_command, out_stream=self.o_stream, err_stream=self.e_stream,
+                                               output_charset=self.output_charset, retry=self.retry, retry_delay=self.retry_delay)
                 self.o_std = response.std_out
                 self.e_std = response.std_err
                 self.stat = response.status_code
 
         except Exception as e:
             self.e_std = e
-            self.stat=-1
+            self.stat = -1
 
 
 def to_text(obj, encoding='utf-8', errors="ignore"):
@@ -181,4 +220,3 @@ def to_bytes(obj, encoding='utf-8', errors="ignore"):
             return obj.encode(encoding, errors)
         except UnicodeEncodeError:
             raise
-

plugin.yaml

@@ -193,6 +193,26 @@ providers:
         renderingOptions:
           groupName: Connection
           instance-scope-node-attribute: "winrm-enable-http-logging"
+      - name: retryconnection
+        title: Retry connection
+        description: "Retry the connection to the node if the connection fails. It can be overwriting at node level using `winrm-retry-connection`"
+        type: Integer
+        default: "1"
+        required: false
+        scope: Instance
+        renderingOptions:
+          groupName: Connection
+          instance-scope-node-attribute: "winrm-retry-connection"
+      - name: retryconnectiondelay
+        title: Retry connection delay
+        description: "Delay between each retry atten (seconds). It can be overwriting at node level using `winrm-retry-connection-delay`"
+        type: Integer
+        default: "10"
+        required: false
+        scope: Instance
+        renderingOptions:
+          groupName: Connection
+          instance-scope-node-attribute: "winrm-retry-connection-delay"
   - name: WinRMcpPython
     title: WinRM Python File Copier
     description: Copying files to remote Windows computer
@@ -352,6 +372,26 @@ providers:
         renderingOptions:
           groupName: Connection
           instance-scope-node-attribute: "winrm-enable-http-logging"
+      - name: retryconnection
+        title: Retry connection
+        description: "Retry the connection to the node if the connection fails. It can be overwriting at node level using `winrm-retry-connection`"
+        type: Integer
+        default: "1"
+        required: false
+        scope: Instance
+        renderingOptions:
+          groupName: Connection
+          instance-scope-node-attribute: "winrm-retry-connection"
+      - name: retryconnectiondelay
+        title: Retry connection delay
+        description: "Delay between each retry atten (seconds). It can be overwriting at node level using `winrm-retry-connection-delay`"
+        type: Integer
+        default: "10"
+        required: false
+        scope: Instance
+        renderingOptions:
+          groupName: Connection
+          instance-scope-node-attribute: "winrm-retry-connection-delay"
   - name: WinRMCheck
     title: WinRM Check Step
     description: Check the connection with a remote node using winrm-python