clean up ssl.py; expose unwrap and add test for it

2 files changed, 29 insertions, 4 deletions

diff --git a/Lib/ssl.py b/Lib/ssl.py
index c072cd9..aa301295 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -75,10 +75,10 @@ from _ssl import (
     SSL_ERROR_INVALID_ERROR_CODE,
     )
 
-from socket import socket, AF_INET, SOCK_STREAM, error
 from socket import getnameinfo as _getnameinfo
 from socket import error as socket_error
 from socket import dup as _dup
+from socket import socket, AF_INET, SOCK_STREAM
 import base64        # for DER-to-PEM translation
 import traceback
 
@@ -296,6 +296,14 @@ class SSLSocket(socket):
         self._sslobj = None
         socket.shutdown(self, how)
 
+    def unwrap (self):
+        if self._sslobj:
+            s = self._sslobj.shutdown()
+            self._sslobj = None
+            return s
+        else:
+            raise ValueError("No SSL wrapper around " + str(self))
+
     def _real_close(self):
         self._sslobj = None
         # self._closed = True
diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py
index 9e36e80..a40a35d 100644
--- a/Lib/test/test_ssl.py
+++ b/Lib/test/test_ssl.py
@@ -279,6 +279,15 @@ else:
                             self.write("OK\n".encode("ASCII", "strict"))
                             if not self.wrap_conn():
                                 return
+                        elif (self.server.starttls_server and self.sslconn
+                              and amsg.strip() == 'ENDTLS'):
+                            if support.verbose and self.server.connectionchatty:
+                                sys.stdout.write(" server: read ENDTLS from client, sending OK...\n")
+                            self.write("OK\n".encode("ASCII", "strict"))
+                            self.sock = self.sslconn.unwrap()
+                            self.sslconn = None
+                            if support.verbose and self.server.connectionchatty:
+                                sys.stdout.write(" server: connection is now unencrypted...\n")
                         else:
                             if (support.verbose and
                                 self.server.connectionchatty):
@@ -868,7 +877,7 @@ else:
 
         def testSTARTTLS (self):
 
-            msgs = ("msg 1", "MSG 2", "STARTTLS", "MSG 3", "msg 4")
+            msgs = ("msg 1", "MSG 2", "STARTTLS", "MSG 3", "msg 4", "ENDTLS", "msg 5", "msg 6")
 
             server = ThreadedEchoServer(CERTFILE,
                                         ssl_version=ssl.PROTOCOL_TLSv1,
@@ -910,8 +919,16 @@ else:
                                     " client:  read %s from server, starting TLS...\n"
                                     % repr(msg))
                             conn = ssl.wrap_socket(s, ssl_version=ssl.PROTOCOL_TLSv1)
-
                             wrapped = True
+                        elif (indata == "ENDTLS" and
+                              str(outdata, 'ASCII', 'replace').strip().lower().startswith("ok")):
+                            if support.verbose:
+                                msg = str(outdata, 'ASCII', 'replace')
+                                sys.stdout.write(
+                                    " client:  read %s from server, ending TLS...\n"
+                                    % repr(msg))
+                            s = conn.unwrap()
+                            wrapped = False
                         else:
                             if support.verbose:
                                 msg = str(outdata, 'ASCII', 'replace')
@@ -922,7 +939,7 @@ else:
                     if wrapped:
                         conn.write("over\n".encode("ASCII", "strict"))
                     else:
-                        s.send("over\n")
+                        s.send("over\n".encode("ASCII", "strict"))
                 if wrapped:
                     conn.close()
                 else: