diff --git a/package/cloudshell/cm/customscript/domain/windows_script_executor.py b/package/cloudshell/cm/customscript/domain/windows_script_executor.py index 0fe2c90..5f75313 100644 --- a/package/cloudshell/cm/customscript/domain/windows_script_executor.py +++ b/package/cloudshell/cm/customscript/domain/windows_script_executor.py @@ -15,6 +15,8 @@ class WindowsScriptExecutor(IScriptExecutor): + COPY_BULK_SIZE = 2000 + def __init__(self, logger, target_host, cancel_sampler): """ :type logger: Logger @@ -41,37 +43,112 @@ def connect(self): except Exception as e: raise ExcutorConnectionError(0, e) + def execute(self, script_file, env_vars, output_writer): + """ + :type script_file: ScriptFile + :type output_writer: ReservationOutputWriter + """ + self.logger.info('Creating temp folder on target machine ...') + tmp_folder = self.create_temp_folder() + self.logger.info('Done (%s).' % tmp_folder) + + try: + self.logger.info('Copying "%s" (%s chars) to "%s" target machine ...' % ( + script_file.name, len(script_file.text), tmp_folder)) + self.copy_script(tmp_folder, script_file) + self.logger.info('Done.') + + self.logger.info('Running "%s" on target machine ...' % script_file.name) + self.run_script(tmp_folder, script_file, env_vars, output_writer) + self.logger.info('Done.') + + finally: + self.logger.info('Deleting "%s" folder from target machine ...' % tmp_folder) + self.delete_temp_folder(tmp_folder) + self.logger.info('Done.') + + def create_temp_folder(self): + """ + :rtype str + """ + code = """ +$fullPath = Join-Path $env:Temp ([System.Guid]::NewGuid().ToString()) +New-Item $fullPath -type directory | Out-Null +Write-Output $fullPath +""" + result = self._run_cancelable(code) + if result.status_code != 0: + raise Exception(ErrorMsg.CREATE_TEMP_FOLDER % result.std_err) + return result.std_out.rstrip('\r\n') + + def copy_script(self, tmp_folder, script_file): + """ + :type tmp_folder: str + :type script_file: ScriptFile + """ + all_size = len(script_file.text) + bulk_zise = WindowsScriptExecutor.COPY_BULK_SIZE + bulks = [script_file.text[i:min(all_size,i+bulk_zise)] for i in range(0, all_size, bulk_zise)] + self.logger.debug("Bulks sizes (%s): %s" % (len(bulks), ', '.join([str(len(b)) for b in bulks]))) + + for bulk in bulks: + encoded_bulk = base64.b64encode(bulk.encode("utf-8")) + code = """ +$path = Join-Path "%s" "%s" +$data = [System.Convert]::FromBase64String("%s") +Add-Content -value $data -encoding byte -path $path +""" + result = self._run_cancelable(code, tmp_folder, script_file.name, encoded_bulk) + if result.status_code != 0: + raise Exception(ErrorMsg.COPY_SCRIPT % result.std_err) + + def run_script(self, tmp_folder, script_file, env_vars, output_writer): """ :type tmp_folder: str :type script_file: ScriptFile :type env_vars: dict :type output_writer: ReservationOutputWriter """ - self.logger.debug('PowerShellScript:' + script_file.text) code = '' for key, value in (env_vars or {}).iteritems(): - code += 'set %s="%s" &&' % (key, str(value)) - encoded_script = base64.b64encode(script_file.text.encode('utf_16_le')).decode('ascii') - code += '\npowershell -encodedcommand %s' % encoded_script - - result = self._run_cancelable(code) + code += '\n$env:%s = "%s"' % (key, str(value)) + code += """ +$path = Join-Path "%s" "%s" +Invoke-Expression $path +""" + result = self._run_cancelable(code, tmp_folder, script_file.name) output_writer.write(result.std_out) output_writer.write(result.std_err) if result.status_code != 0: raise Exception(ErrorMsg.RUN_SCRIPT % result.std_err) + def delete_temp_folder(self, tmp_folder): + """ + :type tmp_folder: str + """ + code = """ +$path = "%s" +Remove-Item $path -recurse +""" + result = self._run_cancelable(code % tmp_folder) + if result.status_code != 0: + raise Exception(ErrorMsg.DELETE_TEMP_FOLDER % result.std_err) - def _run_ps(self, code): - result = self.session.run_ps(code) - self.logger.debug('ReturnedCode:' + str(result.status_code)) - self.logger.debug('Stdout:' + result.std_out) - self.logger.debug('Stderr:' + result.std_err) - return result + # def _run_ps(self, code): + # result = self.session.run_ps(code) + # self.logger.debug('ReturnedCode:' + str(result.status_code)) + # self.logger.debug('Stdout:' + result.std_out) + # self.logger.debug('Stderr:' + result.std_err) + # return result + + def _run_cancelable(self, txt, *args): + ps_code = txt % args + self.logger.debug('PowerShellScript:' + ps_code) - def _run_cancelable(self, code): + bat_code = 'powershell -encodedcommand %s' % base64.b64encode(ps_code.encode('utf_16_le')).decode('ascii') shell_id = self.session.protocol.open_shell() - command_id = self.session.protocol.run_command(shell_id, code) + command_id = self.session.protocol.run_command(shell_id, bat_code) async_result = self.pool.apply_async(self.session.protocol.get_command_output, kwds={'shell_id': shell_id, 'command_id': command_id}) try: diff --git a/package/tests/test_windows_script_executor.py b/package/tests/test_windows_script_executor.py index d86630d..cb84540 100644 --- a/package/tests/test_windows_script_executor.py +++ b/package/tests/test_windows_script_executor.py @@ -24,6 +24,7 @@ def setUp(self): self.session_ctor = self.session_patcher.start() self.session_ctor.return_value = self.session + def tearDown(self): self.session_patcher.stop() @@ -36,27 +37,71 @@ def test_https(self): WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) self.session_ctor.assert_called_with('1.2.3.4', auth=('admin','1234'), transport='ssl') - def test_execute_success(self): + # Create temp folder + + def test_create_temp_folder_success(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('tmp123', '', 0)) + result = executor.create_temp_folder() + self.assertEqual('tmp123', result) + + def test_create_temp_folder_fail(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('', 'some error', 1)) + with self.assertRaises(Exception) as e: + executor.create_temp_folder() + self.assertEqual(ErrorMsg.CREATE_TEMP_FOLDER % 'some error', e.exception.message) + + # Copy script + + def test_copy_script_success(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('','',0)) + executor.copy_script('tmp123', ScriptFile('script1','some script code')) + + def test_copy_long_script_in_bulks(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('','',0)) + executor.copy_script('tmp123', ScriptFile('script1',''.join(['a' for i in range(0,4500)]))) # 3 bulks: 2000,2000,500 + self.assertEqual(3, self.session.protocol.get_command_output.call_count) + + def test_copy_script_fail(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('','some error',1)) + with self.assertRaises(Exception) as e: + executor.copy_script('tmp123', ScriptFile('script1','some script code')) + self.assertEqual(ErrorMsg.COPY_SCRIPT % 'some error', e.exception.message) + + # Run script + + def test_run_script_success(self): executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) output_writer = Mock() self.session.protocol.get_command_output = Mock(return_value=('some output', 'some error', 0)) - executor.execute(ScriptFile('script1', 'some script code'), {'var1':'123'}, output_writer) + executor.run_script('tmp123', ScriptFile('script1', 'some script code'), {'var1':'123'}, output_writer) output_writer.write.assert_any_call('some output') output_writer.write.assert_any_call('some error') - def test_execute_fail(self): + def test_run_script_fail(self): executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) output_writer = Mock() self.session.protocol.get_command_output = Mock(return_value=('some output', 'some error', 1)) with self.assertRaises(Exception, ) as e: - executor.execute(ScriptFile('script1', 'some script code'), {}, output_writer) + executor.run_script('tmp123', ScriptFile('script1', 'some script code'), {}, output_writer) self.assertEqual(ErrorMsg.RUN_SCRIPT % 'some error', e.exception.message) output_writer.write.assert_any_call('some output') output_writer.write.assert_any_call('some error') + # Delete temp folder + + def test_delete_temp_folder_success(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('','',0)) + executor.delete_temp_folder('tmp123') -class Result(object): - def __init__(self, status_code, std_out, std_err): - self.std_err = std_err - self.std_out = std_out - self.status_code = status_code + def test_delete_temp_folder_fail(self): + executor = WindowsScriptExecutor(self.logger, self.host, self.cancel_sampler) + self.session.protocol.get_command_output = Mock(return_value=('','some error',1)) + with self.assertRaises(Exception) as e: + executor.delete_temp_folder('tmp123') + self.assertEqual(ErrorMsg.DELETE_TEMP_FOLDER % 'some error', e.exception.message) \ No newline at end of file diff --git a/package/version.txt b/package/version.txt index afaf360..4575568 100644 --- a/package/version.txt +++ b/package/version.txt @@ -1 +1 @@ -1.0.0 \ No newline at end of file +1.0.0.99 \ No newline at end of file