psrp - use native copy mechanism

changelogs/fragments/psrp-copy.yaml

@@ -0,0 +1,5 @@
+bugfixes:
+- psrp - Use native PSRP mechanism when copying files to support custom endpoints
+
+deprecated_features:
+- psrp - Set the minimum version of ``pypsrp`` to ``0.4.0``.

lib/ansible/plugins/connection/psrp.py

@@ -14,7 +14,7 @@
   underlying transport but instead runs in a PowerShell interpreter.
 version_added: "2.7"
 requirements:
-- pypsrp (Python library)
+- pypsrp>=0.4.0 (Python library)
 options:
   # transport options
   remote_addr:
@@ -311,7 +311,7 @@
 from ansible.plugins.connection import ConnectionBase
 from ansible.plugins.shell.powershell import _common_args
 from ansible.utils.display import Display
-from ansible.utils.hashing import secure_hash
+from ansible.utils.hashing import sha1
 
 HAS_PYPSRP = True
 PYPSRP_IMP_ERR = None
@@ -328,6 +328,12 @@
     HAS_PYPSRP = False
     PYPSRP_IMP_ERR = err
 
+NEWER_PYPSRP = True
+try:
+    import pypsrp.pwsh_scripts
+except ImportError:
+    NEWER_PYPSRP = False
+
 display = Display()
 
 
@@ -446,11 +452,39 @@ def exec_command(self, cmd, in_data=None, sudoable=True):
 
     def put_file(self, in_path, out_path):
         super(Connection, self).put_file(in_path, out_path)
-        display.vvv("PUT %s TO %s" % (in_path, out_path), host=self._psrp_host)
 
         out_path = self._shell._unquote(out_path)
+        display.vvv("PUT %s TO %s" % (in_path, out_path), host=self._psrp_host)
+
+        # The new method that uses PSRP directly relies on a feature added in pypsrp 0.4.0 (release 2019-09-19). In
+        # case someone still has an older version present we warn them asking to update their library to a newer
+        # release and fallback to the old WSMV shell.
+        if NEWER_PYPSRP:
+            rc, stdout, stderr, local_sha1 = self._put_file_new(in_path, out_path)
+
+        else:
+            display.deprecated("Older pypsrp library detected, please update to pypsrp>=0.4.0 to use the newer copy "
+                               "method over PSRP.", version="2.13", collection_name='ansible.builtin')
+            rc, stdout, stderr, local_sha1 = self._put_file_old(in_path, out_path)
+
+        if rc != 0:
+            raise AnsibleError(to_native(stderr))
+
+        put_output = json.loads(stdout)
+        remote_sha1 = put_output.get("sha1")
+
+        if not remote_sha1:
+            raise AnsibleError("Remote sha1 was not returned, stdout: '%s', stderr: '%s'"
+                               % (to_native(stdout), to_native(stderr)))
+
+        if not remote_sha1 == local_sha1:
+            raise AnsibleError("Remote sha1 hash %s does not match local hash %s"
+                               % (to_native(remote_sha1), to_native(local_sha1)))
+
+    def _put_file_old(self, in_path, out_path):
         script = u'''begin {
     $ErrorActionPreference = "Stop"
+    $ProgressPreference = 'SilentlyContinue'
 
     $path = '%s'
     $fd = [System.IO.File]::Create($path)
@@ -467,7 +501,7 @@ def put_file(self, in_path, out_path):
     $hash = $hash.Replace("-", "").ToLowerInvariant()
 
     Write-Output -InputObject "{`"sha1`":`"$hash`"}"
-}''' % self._shell._escape(out_path)
+}''' % out_path
 
         cmd_parts = self._shell._encode_script(script, as_list=True,
                                                strict_mode=False,
@@ -479,6 +513,7 @@ def put_file(self, in_path, out_path):
 
         in_size = os.path.getsize(b_in_path)
         buffer_size = int(self.runspace.connection.max_payload_size / 4 * 3)
+        sha1_hash = sha1()
 
         # copying files is faster when using the raw WinRM shell and not PSRP
         # we will create a WinRS shell just for this process
@@ -496,6 +531,7 @@ def put_file(self, in_path, out_path):
                                   host=self._psrp_host)
                     b64_data = base64.b64encode(data) + b"\r\n"
                     process.send(b64_data, end=(src_file.tell() == in_size))
+                    sha1_hash.update(data)
 
                 # the file was empty, return empty buffer
                 if offset == 0:
@@ -504,22 +540,134 @@ def put_file(self, in_path, out_path):
             process.end_invoke()
             process.signal(SignalCode.CTRL_C)
 
-        if process.rc != 0:
-            raise AnsibleError(to_native(process.stderr))
+        return process.rc, process.stdout, process.stderr, sha1_hash.hexdigest()
 
-        put_output = json.loads(process.stdout)
-        remote_sha1 = put_output.get("sha1")
+    def _put_file_new(self, in_path, out_path):
+        copy_script = '''begin {
+    $ErrorActionPreference = "Stop"
+    $WarningPreference = "Continue"
+    $path = $MyInvocation.UnboundArguments[0]
+    $fd = [System.IO.File]::Create($path)
+    $algo = [System.Security.Cryptography.SHA1CryptoServiceProvider]::Create()
+    $bytes = @()
 
-        if not remote_sha1:
-            raise AnsibleError("Remote sha1 was not returned, stdout: '%s', "
-                               "stderr: '%s'" % (to_native(process.stdout),
-                                                 to_native(process.stderr)))
+    $bindingFlags = [System.Reflection.BindingFlags]'NonPublic, Instance'
+    Function Get-Property {
+        <#
+        .SYNOPSIS
+        Gets the private/internal property specified of the object passed in.
+        #>
+        Param (
+            [Parameter(Mandatory=$true, ValueFromPipeline=$true)]
+            [System.Object]
+            $Object,
+
+            [Parameter(Mandatory=$true, Position=1)]
+            [System.String]
+            $Name
+        )
 
-        local_sha1 = secure_hash(in_path)
-        if not remote_sha1 == local_sha1:
-            raise AnsibleError("Remote sha1 hash %s does not match local hash "
-                               "%s" % (to_native(remote_sha1),
-                                       to_native(local_sha1)))
+        $Object.GetType().GetProperty($Name, $bindingFlags).GetValue($Object, $null)
+    }
+
+    Function Set-Property {
+        <#
+        .SYNOPSIS
+        Sets the private/internal property specified on the object passed in.
+        #>
+        Param (
+            [Parameter(Mandatory=$true, ValueFromPipeline=$true)]
+            [System.Object]
+            $Object,
+
+            [Parameter(Mandatory=$true, Position=1)]
+            [System.String]
+            $Name,
+
+            [Parameter(Mandatory=$true, Position=2)]
+            [AllowNull()]
+            [System.Object]
+            $Value
+        )
+
+        $Object.GetType().GetProperty($Name, $bindingFlags).SetValue($Object, $Value, $null)
+    }
+
+    Function Get-Field {
+        <#
+        .SYNOPSIS
+        Gets the private/internal field specified of the object passed in.
+        #>
+        Param (
+            [Parameter(Mandatory=$true, ValueFromPipeline=$true)]
+            [System.Object]
+            $Object,
+
+            [Parameter(Mandatory=$true, Position=1)]
+            [System.String]
+            $Name
+        )
+
+        $Object.GetType().GetField($Name, $bindingFlags).GetValue($Object)
+    }
+
+    # MaximumAllowedMemory is required to be set to so we can send input data that exceeds the limit on a PS
+    # Runspace. We use reflection to access/set this property as it is not accessible publicly. This is not ideal
+    # but works on all PowerShell versions I've tested with. We originally used WinRS to send the raw bytes to the
+    # host but this falls flat if someone is using a custom PS configuration name so this is a workaround. This
+    # isn't required for smaller files so if it fails we ignore the error and hope it wasn't needed.
+    # https://github.com/PowerShell/PowerShell/blob/c8e72d1e664b1ee04a14f226adf655cced24e5f0/src/System.Management.Automation/engine/serialization.cs#L325
+    try {
+        $Host | Get-Property 'ExternalHost' | `
+            Get-Field '_transportManager' | `
+            Get-Property 'Fragmentor' | `
+            Get-Property 'DeserializationContext' | `
+            Set-Property 'MaximumAllowedMemory' $null
+    } catch {}
+}
+process {
+    $bytes = [System.Convert]::FromBase64String($input)
+    $algo.TransformBlock($bytes, 0, $bytes.Length, $bytes, 0) > $null
+    $fd.Write($bytes, 0, $bytes.Length)
+}
+end {
+    $fd.Close()
+
+    $algo.TransformFinalBlock($bytes, 0, 0) > $null
+    $hash = [System.BitConverter]::ToString($algo.Hash).Replace('-', '').ToLowerInvariant()
+    Write-Output -InputObject "{`"sha1`":`"$hash`"}"
+}
+'''
+
+        # Get the buffer size of each fragment to send, subtract 82 for the fragment, message, and other header info
+        # fields that PSRP adds. Adjust to size of the base64 encoded bytes length.
+        buffer_size = int((self.runspace.connection.max_payload_size - 82) / 4 * 3)
+
+        sha1_hash = sha1()
+
+        b_in_path = to_bytes(in_path, errors='surrogate_or_strict')
+        if not os.path.exists(b_in_path):
+            raise AnsibleFileNotFound('file or module does not exist: "%s"' % to_native(in_path))
+
+        def read_gen():
+            offset = 0
+
+            with open(b_in_path, 'rb') as src_fd:
+                for b_data in iter((lambda: src_fd.read(buffer_size)), b""):
+                    data_len = len(b_data)
+                    offset += data_len
+                    sha1_hash.update(b_data)
+
+                    # PSRP technically supports sending raw bytes but that method requires a larger CLIXML message.
+                    # Sending base64 is still more efficient here.
+                    display.vvvvv("PSRP PUT %s to %s (offset=%d, size=%d" % (in_path, out_path, offset, data_len),
+                                  host=self._psrp_host)
+                    b64_data = base64.b64encode(b_data)
+                    yield [to_text(b64_data)]
+
+        rc, stdout, stderr = self._exec_psrp_script(copy_script, read_gen(), arguments=[out_path], force_stop=True)
+
+        return rc, stdout, stderr, sha1_hash.hexdigest()
 
     def fetch_file(self, in_path, out_path):
         super(Connection, self).fetch_file(in_path, out_path)
@@ -713,9 +861,13 @@ def _build_kwargs(self):
             option = self.get_option('_extras')['ansible_psrp_%s' % arg]
             self._psrp_conn_kwargs[arg] = option
 
-    def _exec_psrp_script(self, script, input_data=None, use_local_scope=True, force_stop=False):
+    def _exec_psrp_script(self, script, input_data=None, use_local_scope=True, force_stop=False, arguments=None):
         ps = PowerShell(self.runspace)
         ps.add_script(script, use_local_scope=use_local_scope)
+        if arguments:
+            for arg in arguments:
+                ps.add_argument(arg)
+
         ps.invoke(input=input_data)
 
         rc, stdout, stderr = self._parse_pipeline_result(ps)