Copy win ps script (again) + Do it in bulks

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:

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)

package/version.txt

@@ -1 +1 @@
-1.0.0
+1.0.0.99