From 549163170fe34c3f1bc066ca69137522133b16b1 Mon Sep 17 00:00:00 2001 From: nitzmahone Date: Fri, 23 Oct 2015 14:37:50 -0700 Subject: [PATCH] fast winrm put_file without size restrictions --- lib/ansible/plugins/connection/winrm.py | 137 ++++++++++++++++++------ 1 file changed, 102 insertions(+), 35 deletions(-) diff --git a/lib/ansible/plugins/connection/winrm.py b/lib/ansible/plugins/connection/winrm.py index 5ae83adab3..5e75da1dc9 100644 --- a/lib/ansible/plugins/connection/winrm.py +++ b/lib/ansible/plugins/connection/winrm.py @@ -24,6 +24,8 @@ import os import re import shlex import traceback +import json +import xmltodict from ansible.compat.six.moves.urllib.parse import urlunsplit @@ -44,6 +46,7 @@ except ImportError: from ansible.errors import AnsibleFileNotFound from ansible.plugins.connection import ConnectionBase +from ansible.utils.hashing import secure_hash from ansible.utils.path import makedirs_safe from ansible.utils.unicode import to_bytes, to_unicode, to_str from ansible.utils.vars import combine_vars @@ -151,7 +154,21 @@ class Connection(ConnectionBase): else: raise AnsibleError('No transport found for WinRM connection') - def _winrm_exec(self, command, args=(), from_exec=False): + def _winrm_send_input(self, protocol, shell_id, command_id, stdin, eof=False): + rq = {'env:Envelope': protocol._get_soap_header( + resource_uri='http://schemas.microsoft.com/wbem/wsman/1/windows/shell/cmd', + action='http://schemas.microsoft.com/wbem/wsman/1/windows/shell/Send', + shell_id=shell_id)} + stream = rq['env:Envelope'].setdefault('env:Body', {}).setdefault('rsp:Send', {})\ + .setdefault('rsp:Stream', {}) + stream['@Name'] = 'stdin' + stream['@CommandId'] = command_id + stream['#text'] = base64.b64encode(to_bytes(stdin)) + if eof: + stream['@End'] = 'true' + rs = protocol.send_message(xmltodict.unparse(rq)) + + def _winrm_exec(self, command, args=(), from_exec=False, stdin_iterator=None): if from_exec: display.vvvvv("WINRM EXEC %r %r" % (command, args), host=self._winrm_host) else: @@ -162,7 +179,19 @@ class Connection(ConnectionBase): self.shell_id = self.protocol.open_shell(codepage=65001) # UTF-8 command_id = None try: - command_id = self.protocol.run_command(self.shell_id, to_bytes(command), map(to_bytes, args)) + command_id = self.protocol.run_command(self.shell_id, to_bytes(command), map(to_bytes, args), console_mode_stdin=(stdin_iterator == None)) + + # TODO: try/except around this, so we can get/return the command result on a broken pipe or other failure (probably more useful than the 500 that comes from this) + try: + if stdin_iterator: + for (data, is_last) in stdin_iterator: + self._winrm_send_input(self.protocol, self.shell_id, command_id, data, eof=is_last) + except: + # TODO: set/propagate an error flag, but don't throw (or include the command output in the exception) + pass + + # NB: this could hang if the receiver is still running (eg, network failed a Send request but the server's still happy). + # Consider adding pywinrm status check/abort operations to see if the target is still running after a failure. response = Response(self.protocol.get_command_output(self.shell_id, command_id)) if from_exec: display.vvvvv('WINRM RESULT %r' % to_unicode(response), host=self._winrm_host) @@ -212,45 +241,83 @@ class Connection(ConnectionBase): result.std_err = to_bytes(result.std_err) return (result.status_code, result.std_out, result.std_err) + # FUTURE: determine buffer size at runtime via remote winrm config? + def _put_file_stdin_iterator(self, in_path, out_path, buffer_size=250000): + in_size = os.path.getsize(in_path) + offset = 0 + with open(in_path, 'rb') as in_file: + for out_data in iter((lambda:in_file.read(buffer_size)), ''): + offset += len(out_data) + self._display.vvvvv('WINRM PUT "%s" to "%s" (offset=%d size=%d)' % (in_path, out_path, offset, len(out_data)), host=self._winrm_host) + # yes, we're double-encoding over the wire in this case- we want to ensure that the data shipped to the end PS pipeline is still b64-encoded + b64_data = base64.b64encode(out_data) + '\r\n' + # cough up the data, as well as an indicator if this is the last chunk so winrm_send knows to set the End signal + yield b64_data, (in_file.tell() == in_size) + + if offset == 0: # empty file, return an empty buffer + eof to close it + yield "", True + def put_file(self, in_path, out_path): super(Connection, self).put_file(in_path, out_path) out_path = self._shell._unquote(out_path) display.vvv('PUT "%s" TO "%s"' % (in_path, out_path), host=self._winrm_host) if not os.path.exists(in_path): raise AnsibleFileNotFound('file or module does not exist: "%s"' % in_path) - with open(in_path) as in_file: - in_size = os.path.getsize(in_path) - script_template = ''' - $s = [System.IO.File]::OpenWrite("%s"); - [void]$s.Seek(%d, [System.IO.SeekOrigin]::Begin); - $b = [System.Convert]::FromBase64String("%s"); - [void]$s.Write($b, 0, $b.length); - [void]$s.SetLength(%d); - [void]$s.Close(); - ''' - # Determine max size of data we can pass per command. - script = script_template % (self._shell._escape(out_path), in_size, '', in_size) - cmd = self._shell._encode_script(script) - # Encode script with no data, subtract its length from 8190 (max - # windows command length), divide by 2.67 (UTF16LE base64 command - # encoding), then by 1.35 again (data base64 encoding). - buffer_size = int(((8190 - len(cmd)) / 2.67) / 1.35) - for offset in xrange(0, in_size or 1, buffer_size): - try: - out_data = in_file.read(buffer_size) - if offset == 0: - if out_data.lower().startswith('#!powershell') and not out_path.lower().endswith('.ps1'): - out_path = out_path + '.ps1' - b64_data = base64.b64encode(out_data) - script = script_template % (self._shell._escape(out_path), offset, b64_data, in_size) - display.vvvvv('WINRM PUT "%s" to "%s" (offset=%d size=%d)' % (in_path, out_path, offset, len(out_data)), host=self._winrm_host) - cmd_parts = self._shell._encode_script(script, as_list=True) - result = self._winrm_exec(cmd_parts[0], cmd_parts[1:]) - if result.status_code != 0: - raise IOError(to_str(result.std_err)) - except Exception: - traceback.print_exc() - raise AnsibleError('failed to transfer file to "%s"' % out_path) + + script_template = ''' + begin {{ + $path = "{0}" + + $DebugPreference = "Continue" + $ErrorActionPreference = "Stop" + Set-StrictMode -Version 2 + + $fd = [System.IO.File]::Create($path) + + $sha1 = [System.Security.Cryptography.SHA1CryptoServiceProvider]::Create() + + $bytes = @() #initialize for empty file case + }} + process {{ + $bytes = [System.Convert]::FromBase64String($input) + $sha1.TransformBlock($bytes, 0, $bytes.Length, $bytes, 0) | Out-Null + $fd.Write($bytes, 0, $bytes.Length) + }} + end {{ + $sha1.TransformFinalBlock($bytes, 0, 0) | Out-Null + + $hash = [System.BitConverter]::ToString($sha1.Hash).Replace("-", "").ToLowerInvariant() + + $fd.Close() + + Write-Output "{{""sha1"":""$hash""}}" + }} + ''' + + # FUTURE: this sucks- why can't the module/shell stuff do this? + with open(in_path, 'r') as temp_file: + if temp_file.read(15).lower().startswith('#!powershell') and not out_path.lower().endswith('.ps1'): + out_path = out_path + '.ps1' + + script = script_template.format(self._shell._escape(out_path)) + cmd_parts = self._shell._encode_script(script, as_list=True, strict_mode=False) + + result = self._winrm_exec(cmd_parts[0], cmd_parts[1:], stdin_iterator=self._put_file_stdin_iterator(in_path, out_path)) + # TODO: improve error handling + if result.status_code != 0: + raise IOError(to_str(result.std_err)) + + put_output = json.loads(result.std_out) + remote_sha1 = put_output.get("sha1") + + if not remote_sha1: + raise IOError("Remote sha1 was not returned") + + local_sha1 = secure_hash(in_path) + + if not remote_sha1 == local_sha1: + raise IOError("Remote sha1 hash {0} does not match local hash {1}".format(remote_sha1, local_sha1)) + def fetch_file(self, in_path, out_path): super(Connection, self).fetch_file(in_path, out_path)