Merge 2528b85 into 259dd91

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Added `proxy` support where it can be defined within the application, with the ability to specify the proxy within the application
 - Fix for shell not setting all environment variables.
 - Fix session clixml encoding on Python 3
+- `Protocol.close_shell(shell_id)` will now close the session(and TCP connections) to the Windows machine. `close_session` option has been added in case of leaving the session alone.
 
 ### Version 0.3.0
 - Added support for message encryption over HTTP when using NTLM/Kerberos/CredSSP

winrm/protocol.py

@@ -277,30 +277,39 @@ def send_message(self, message):
 
                 raise WinRMError('{0} (extended fault data: {1})'.format(error_message, fault_data))
 
-    def close_shell(self, shell_id):
+    def close_shell(self, shell_id, close_session=True):
         """
         Close the shell
         @param string shell_id: The shell id on the remote machine.
          See #open_shell
+        @param bool close_session: If we want to close the requests's session.
+         Allows to completely close all TCP connections to the server.
         @returns This should have more error checking but it just returns true
          for now.
         @rtype bool
         """
-        message_id = uuid.uuid4()
-        req = {'env:Envelope': self._get_soap_header(
-            resource_uri='http://schemas.microsoft.com/wbem/wsman/1/windows/shell/cmd',  # NOQA
-            action='http://schemas.xmlsoap.org/ws/2004/09/transfer/Delete',
-            shell_id=shell_id,
-            message_id=message_id)}
-
-        # SOAP message requires empty env:Body
-        req['env:Envelope'].setdefault('env:Body', {})
+        try:
+            message_id = uuid.uuid4()
+            req = {'env:Envelope': self._get_soap_header(
+                resource_uri='http://schemas.microsoft.com/wbem/wsman/1/windows/shell/cmd',  # NOQA
+                action='http://schemas.xmlsoap.org/ws/2004/09/transfer/Delete',
+                shell_id=shell_id,
+                message_id=message_id)}
+
+            # SOAP message requires empty env:Body
+            req['env:Envelope'].setdefault('env:Body', {})
+
+            res = self.send_message(xmltodict.unparse(req))
+            root = ET.fromstring(res)
+            relates_to = next(
+                node for node in root.findall('.//*')
+                if node.tag.endswith('RelatesTo')).text
+        finally:
+            # Close the transport if we are done with the shell.
+            # This will ensure no lingering TCP connections are thrown back into a requests' connection pool.
+            if close_session:
+                self.transport.close_session()
 
-        res = self.send_message(xmltodict.unparse(req))
-        root = ET.fromstring(res)
-        relates_to = next(
-            node for node in root.findall('.//*')
-            if node.tag.endswith('RelatesTo')).text
         # TODO change assert into user-friendly exception
         assert uuid.UUID(relates_to.replace('uuid:', '')) == message_id
 

winrm/tests/conftest.py

@@ -351,6 +351,9 @@ def send_message(self, message):
         else:
             raise Exception('Message was not expected\n\n%s' % message)
 
+    def close_session(self):
+        pass
+
 
 @fixture(scope='module')
 def protocol_fake(request):

winrm/tests/test_protocol.py

@@ -7,7 +7,7 @@ def test_open_shell_and_close_shell(protocol_fake):
     shell_id = protocol_fake.open_shell()
     assert shell_id == '11111111-1111-1111-1111-111111111113'
 
-    protocol_fake.close_shell(shell_id)
+    protocol_fake.close_shell(shell_id, close_session=True)
 
 
 def test_run_command_with_arguments_and_cleanup_command(protocol_fake):

winrm/tests/test_transport.py

@@ -1,5 +1,6 @@
 # coding=utf-8
 import os
+import mock
 import unittest
 
 from winrm import transport
@@ -248,3 +249,28 @@ def test_build_session_invalid_encryption(self):
                                 message_encryption='invalid_value'
                                 )
         self.assertEqual("invalid message_encryption arg: invalid_value. Should be 'auto', 'always', or 'never'", str(exc.exception))
+
+    @mock.patch('requests.Session')
+    def test_close_session(self, mock_session):
+        t_default = transport.Transport(endpoint="Endpoint",
+                                        server_cert_validation='ignore',
+                                        username='test',
+                                        password='test',
+                                        auth_method='basic',
+                                        )
+        t_default.build_session()
+        t_default.close_session()
+        mock_session.return_value.close.assert_called_once_with()
+        self.assertIsNone(t_default.session)
+
+    @mock.patch('requests.Session')
+    def test_close_session_not_built(self, mock_session):
+        t_default = transport.Transport(endpoint="Endpoint",
+                                        server_cert_validation='ignore',
+                                        username='test',
+                                        password='test',
+                                        auth_method='basic',
+                                        )
+        t_default.close_session()
+        self.assertFalse(mock_session.return_value.close.called)
+        self.assertIsNone(t_default.session)

winrm/transport.py

@@ -278,6 +278,12 @@ def setup_encryption(self):
         self._send_message_request(prepared_request, '')
         self.encryption = Encryption(self.session, self.auth_method)
 
+    def close_session(self):
+        if not self.session:
+            return
+        self.session.close()
+        self.session = None
+
     def send_message(self, message):
         if not self.session:
             self.build_session()