# Test the support for SSL and sockets

import sys

import unittest

from test import test_support

import asyncore

import socket

import select

import errno

import subprocess

import time

import os

import pprint

import urllib, urlparse

import shutil

import traceback

from BaseHTTPServer import HTTPServer

from SimpleHTTPServer import SimpleHTTPRequestHandler

# Optionally test SSL support, if we have it in the tested platform

skip_expected = False

try:

    import ssl

except ImportError:

    skip_expected = True

HOST = test_support.HOST

CERTFILE = None

SVN_PYTHON_ORG_ROOT_CERT = None

def handle_error(prefix):

    exc_format = ' '.join(traceback.format_exception(*sys.exc_info()))

    if test_support.verbose:

        sys.stdout.write(prefix + exc_format)

    def testSimpleSSLwrap(self):

        try:

            ssl.sslwrap_simple(socket.socket(socket.AF_INET))

        except IOError, e:

            if e.errno == 32: # broken pipe when ssl_sock.do_handshake(), this test doesn't care about that

                pass

            else:

                raise

        try:

            ssl.sslwrap_simple(socket.socket(socket.AF_INET)._sock)

        except IOError, e:

            if e.errno == 32: # broken pipe when ssl_sock.do_handshake(), this test doesn't care about that

                pass

            else:

                raise

class BasicTests(unittest.TestCase):

    def testSSLconnect(self):

        if not test_support.is_resource_enabled('network'):

            return

        s = ssl.wrap_socket(socket.socket(socket.AF_INET),

                            cert_reqs=ssl.CERT_NONE)

        s.connect(("svn.python.org", 443))

        c = s.getpeercert()

        if c:

            raise test_support.TestFailed("Peer cert %s shouldn't be here!")

        s.close()

        # this should fail because we have no verification certs

        s = ssl.wrap_socket(socket.socket(socket.AF_INET),

                            cert_reqs=ssl.CERT_REQUIRED)

        try:

            s.connect(("svn.python.org", 443))

        except ssl.SSLError:

            pass

        finally:

            s.close()

    def testCrucialConstants(self):

        ssl.PROTOCOL_SSLv2

        ssl.PROTOCOL_SSLv23

        ssl.PROTOCOL_SSLv3

        ssl.PROTOCOL_TLSv1

        ssl.CERT_NONE

        ssl.CERT_OPTIONAL

        ssl.CERT_REQUIRED

    def testRAND(self):

        v = ssl.RAND_status()

        if test_support.verbose:

            sys.stdout.write("\n RAND_status is %d (%s)\n"

                             % (v, (v and "sufficient randomness") or

                                "insufficient randomness"))

        try:

            ssl.RAND_egd(1)

        except TypeError:

            pass

        else:

            print "didn't raise TypeError"

        ssl.RAND_add("this is a random string", 75.0)

    def testParseCert(self):

        # note that this uses an 'unofficial' function in _ssl.c,

        # provided solely for this test, to exercise the certificate

        # parsing code

        p = ssl._ssl._test_decode_cert(CERTFILE, False)

        if test_support.verbose:

            sys.stdout.write("\n" + pprint.pformat(p) + "\n")

    def testDERtoPEM(self):

        pem = open(SVN_PYTHON_ORG_ROOT_CERT, 'r').read()

        d1 = ssl.PEM_cert_to_DER_cert(pem)

        p2 = ssl.DER_cert_to_PEM_cert(d1)

        d2 = ssl.PEM_cert_to_DER_cert(p2)

        if (d1 != d2):

            raise test_support.TestFailed("PEM-to-DER or DER-to-PEM translation failed")

class NetworkedTests(unittest.TestCase):

    def testConnect(self):

        s = ssl.wrap_socket(socket.socket(socket.AF_INET),

                            cert_reqs=ssl.CERT_NONE)

        s.connect(("svn.python.org", 443))

        c = s.getpeercert()

        if c:

            raise test_support.TestFailed("Peer cert %s shouldn't be here!")

        s.close()

        # this should fail because we have no verification certs

        s = ssl.wrap_socket(socket.socket(socket.AF_INET),

                            cert_reqs=ssl.CERT_REQUIRED)

        try:

            s.connect(("svn.python.org", 443))

        except ssl.SSLError:

            pass

        finally:

            s.close()

        # this should succeed because we specify the root cert

        s = ssl.wrap_socket(socket.socket(socket.AF_INET),

                            cert_reqs=ssl.CERT_REQUIRED,

                            ca_certs=SVN_PYTHON_ORG_ROOT_CERT)

        try:

            s.connect(("svn.python.org", 443))

        except ssl.SSLError, x:

            raise test_support.TestFailed("Unexpected exception %s" % x)

        finally:

            s.close()

    def testNonBlockingHandshake(self):

        s = socket.socket(socket.AF_INET)

        s.connect(("svn.python.org", 443))

        s.setblocking(False)

        s = ssl.wrap_socket(s,

                            cert_reqs=ssl.CERT_NONE,

                            do_handshake_on_connect=False)

        count = 0

        while True:

            try:

                count += 1

                s.do_handshake()

                break

            except ssl.SSLError, err:

                if err.args[0] == ssl.SSL_ERROR_WANT_READ:

                    select.select([s], [], [])

                elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:

                    select.select([], [s], [])

                else:

                    raise

        s.close()

        if test_support.verbose:

            sys.stdout.write("\nNeeded %d calls to do_handshake() to establish session.\n" % count)

    def testFetchServerCert(self):

        pem = ssl.get_server_certificate(("svn.python.org", 443))

        if not pem:

            raise test_support.TestFailed("No server certificate on svn.python.org:443!")

        try:

            pem = ssl.get_server_certificate(("svn.python.org", 443), ca_certs=CERTFILE)

        except ssl.SSLError:

            #should fail

            pass

        else:

            raise test_support.TestFailed("Got server certificate %s for svn.python.org!" % pem)

        pem = ssl.get_server_certificate(("svn.python.org", 443), ca_certs=SVN_PYTHON_ORG_ROOT_CERT)

        if not pem:

            raise test_support.TestFailed("No server certificate on svn.python.org:443!")

        if test_support.verbose:

            sys.stdout.write("\nVerified certificate for svn.python.org:443 is\n%s\n" % pem)

try:

    import threading

except ImportError:

    _have_threads = False

else:

    _have_threads = True

    class ThreadedEchoServer(threading.Thread):

        class ConnectionHandler(threading.Thread):

            """A mildly complicated class, because we want it to work both

            with and without the SSL wrapper around the socket connection, so

            that we can test the STARTTLS functionality."""

            def __init__(self, server, connsock):

                self.server = server

                self.running = False

                self.sock = connsock

                self.sock.setblocking(1)

                self.sslconn = None

                threading.Thread.__init__(self)

                self.daemon = True

            def show_conn_details(self):

                if self.server.certreqs == ssl.CERT_REQUIRED:

                    cert = self.sslconn.getpeercert()

                    if test_support.verbose and self.server.chatty:

                        sys.stdout.write(" client cert is " + pprint.pformat(cert) + "\n")

                    cert_binary = self.sslconn.getpeercert(True)

                    if test_support.verbose and self.server.chatty:

                        sys.stdout.write(" cert binary is " + str(len(cert_binary)) + " bytes\n")

                cipher = self.sslconn.cipher()

                if test_support.verbose and self.server.chatty:

                    sys.stdout.write(" server: connection cipher is now " + str(cipher) + "\n")

            def wrap_conn (self):

                try:

                    self.sslconn = ssl.wrap_socket(self.sock, server_side=True,

                                                   certfile=self.server.certificate,

                                                   ssl_version=self.server.protocol,

                                                   ca_certs=self.server.cacerts,

                                                   cert_reqs=self.server.certreqs)

                except:

                    if self.server.chatty:

                        handle_error("\n server:  bad connection attempt from " +

                                     str(self.sock.getpeername()) + ":\n")

                    self.close()

                    if not self.server.expect_bad_connects:

                        # here, we want to stop the server, because this shouldn't

                        # happen in the context of our test case

                        self.running = False

                        # normally, we'd just stop here, but for the test

                        # harness, we want to stop the server

                        self.server.stop()

                    return False

                else:

                    return True

            def read(self):

                if self.sslconn:

                    return self.sslconn.read()

                else:

                    return self.sock.recv(1024)

            def write(self, bytes):

                if self.sslconn:

                    return self.sslconn.write(bytes)

                else:

                    return self.sock.send(bytes)

            def close(self):

                if self.sslconn:

                    self.sslconn.close()

                else:

                    self.sock._sock.close()

            def run (self):

                self.running = True

                if not self.server.starttls_server:

                    if isinstance(self.sock, ssl.SSLSocket):

                        self.sslconn = self.sock

                    elif not self.wrap_conn():

                        return

                    self.show_conn_details()

                while self.running:

                    try:

                        msg = self.read()

                        if not msg:

                            # eof, so quit this handler

                            self.running = False

                            self.close()

                        elif msg.strip() == 'over':

                            if test_support.verbose and self.server.connectionchatty:

                                sys.stdout.write(" server: client closed connection\n")

                            self.close()

                            return

                        elif self.server.starttls_server and msg.strip() == 'STARTTLS':

                            if test_support.verbose and self.server.connectionchatty:

                                sys.stdout.write(" server: read STARTTLS from client, sending OK...\n")

                            self.write("OK\n")

                            if not self.wrap_conn():

                                return

                        elif self.server.starttls_server and self.sslconn and msg.strip() == 'ENDTLS':

                            if test_support.verbose and self.server.connectionchatty:

                                sys.stdout.write(" server: read ENDTLS from client, sending OK...\n")

                            self.write("OK\n")

                            self.sslconn.unwrap()

                            self.sslconn = None

                            if test_support.verbose and self.server.connectionchatty:

                                sys.stdout.write(" server: connection is now unencrypted...\n")

                        else:

                            if (test_support.verbose and

                                self.server.connectionchatty):

                                ctype = (self.sslconn and "encrypted") or "unencrypted"

                                sys.stdout.write(" server: read %s (%s), sending back %s (%s)...\n"

                                                 % (repr(msg), ctype, repr(msg.lower()), ctype))

                            self.write(msg.lower())

                    except ssl.SSLError:

                        if self.server.chatty:

                            handle_error("Test server failure:\n")

                        self.close()

                        self.running = False

                        # normally, we'd just stop here, but for the test

                        # harness, we want to stop the server

                        self.server.stop()

                    except:

                        handle_error('')

        def __init__(self, certificate, ssl_version=None,

                     certreqs=None, cacerts=None, expect_bad_connects=False,

                     chatty=True, connectionchatty=False, starttls_server=False,

                     wrap_accepting_socket=False):

            if ssl_version is None:

                ssl_version = ssl.PROTOCOL_TLSv1

            if certreqs is None:

                certreqs = ssl.CERT_NONE

            self.certificate = certificate

            self.protocol = ssl_version

            self.certreqs = certreqs

            self.cacerts = cacerts

            self.expect_bad_connects = expect_bad_connects

            self.chatty = chatty

            self.connectionchatty = connectionchatty

            self.starttls_server = starttls_server

            self.sock = socket.socket()

            self.flag = None

            if wrap_accepting_socket:

                self.sock = ssl.wrap_socket(self.sock, server_side=True,

                                            certfile=self.certificate,

                                            cert_reqs = self.certreqs,

                                            ca_certs = self.cacerts,

                                            ssl_version = self.protocol)

                if test_support.verbose and self.chatty:

                    sys.stdout.write(' server:  wrapped server socket as %s\n' % str(self.sock))

            self.port = test_support.bind_port(self.sock)

            self.active = False

            threading.Thread.__init__(self)

            self.daemon = True

        def start (self, flag=None):

            self.flag = flag

            threading.Thread.start(self)

        def run (self):

            self.sock.settimeout(0.5)

            self.sock.listen(5)

            self.active = True

            if self.flag:

                # signal an event

                self.flag.set()

            while self.active:

                try:

                    newconn, connaddr = self.sock.accept()

                    if test_support.verbose and self.chatty:

                        sys.stdout.write(' server:  new connection from '

                                         + str(connaddr) + '\n')

                    handler = self.ConnectionHandler(self, newconn)

                    handler.start()

                except socket.timeout:

                    pass

                except KeyboardInterrupt:

                    self.stop()

                except:

                    if self.chatty:

                        handle_error("Test server failure:\n")

            self.sock.close()

        def stop (self):

            self.active = False

    class AsyncoreEchoServer(threading.Thread):

        class EchoServer (asyncore.dispatcher):

            class ConnectionHandler (asyncore.dispatcher_with_send):

                def __init__(self, conn, certfile):

                    asyncore.dispatcher_with_send.__init__(self, conn)

                    self.socket = ssl.wrap_socket(conn, server_side=True,

                                                  certfile=certfile,

                                                  do_handshake_on_connect=True)

                def readable(self):

                    if isinstance(self.socket, ssl.SSLSocket):

                        while self.socket.pending() > 0:

                            self.handle_read_event()

                    return True

                def handle_read(self):

                    data = self.recv(1024)

                    self.send(data.lower())

                def handle_close(self):

                    self.close()

                    if test_support.verbose:

                        sys.stdout.write(" server:  closed connection %s\n" % self.socket)

                def handle_error(self):

                    raise

            def __init__(self, certfile):

                self.certfile = certfile

                asyncore.dispatcher.__init__(self)

                self.create_socket(socket.AF_INET, socket.SOCK_STREAM)

                self.port = test_support.bind_port(self.socket)

                self.listen(5)

            def handle_accept(self):

                sock_obj, addr = self.accept()

                if test_support.verbose:

                    sys.stdout.write(" server:  new connection from %s:%s\n" %addr)

                self.ConnectionHandler(sock_obj, self.certfile)

            def handle_error(self):

                raise

        def __init__(self, certfile):

            self.flag = None

            self.active = False

            self.server = self.EchoServer(certfile)

            self.port = self.server.port

            threading.Thread.__init__(self)

            self.daemon = True

        def __str__(self):

            return "<%s %s>" % (self.__class__.__name__, self.server)

        def start (self, flag=None):

            self.flag = flag

            threading.Thread.start(self)

        def run (self):

            self.active = True

            if self.flag:

                self.flag.set()

            while self.active:

                try:

                    asyncore.loop(1)

                except:

                    pass

        def stop (self):

            self.active = False

            self.server.close()

    class SocketServerHTTPSServer(threading.Thread):

        class HTTPSServer(HTTPServer):

            def __init__(self, server_address, RequestHandlerClass, certfile):

                HTTPServer.__init__(self, server_address, RequestHandlerClass)

                # we assume the certfile contains both private key and certificate

                self.certfile = certfile

                self.active = False

                self.active_lock = threading.Lock()

                self.allow_reuse_address = True

            def __str__(self):

                return ('<%s %s:%s>' %

                        (self.__class__.__name__,

                         self.server_name,

                         self.server_port))

            def get_request (self):

                # override this to wrap socket with SSL

                sock, addr = self.socket.accept()

                sslconn = ssl.wrap_socket(sock, server_side=True,

                                          certfile=self.certfile)

                return sslconn, addr

            # The methods overridden below this are mainly so that we

            # can run it in a thread and be able to stop it from another

            # You probably wouldn't need them in other uses.

            def server_activate(self):

                # We want to run this in a thread for testing purposes,

                # so we override this to set timeout, so that we get

                # a chance to stop the server

                self.socket.settimeout(0.5)

                HTTPServer.server_activate(self)

            def serve_forever(self):

                # We want this to run in a thread, so we use a slightly

                # modified version of "forever".

                self.active = True

                while 1:

                    try:

                        # We need to lock while handling the request.

                        # Another thread can close the socket after self.active

                        # has been checked and before the request is handled.

                        # This causes an exception when using the closed socket.

                        with self.active_lock:

                            if not self.active:

                                break

                            self.handle_request()

                    except socket.timeout:

                        pass

                    except KeyboardInterrupt:

                        self.server_close()

                        return

                    except:

                        sys.stdout.write(''.join(traceback.format_exception(*sys.exc_info())))

                        break

                    time.sleep(0.1)

            def server_close(self):

                # Again, we want this to run in a thread, so we need to override

                # close to clear the "active" flag, so that serve_forever() will

                # terminate.

                with self.active_lock:

                    HTTPServer.server_close(self)