SSL Socket Context Manager

sslwrap.py

#!/usr/bin/env python
__author__ = 'Charles Hooper <[email protected]>'

import socket, ssl
socket.setdefaulttimeout(3)

class EasySSLWrap(object):
    """Connect over TCP+SSL using a context manager.

    Usage:

    >>> addr = socket.gethostbyname('www.google.com')
    >>> port = 443
    >>> with EasySSLWrap( addr, port ) as sslsock:
    >>>     print sslsock.getpeercert()"""

    socket_family = socket.AF_INET
    socket_type = socket.SOCK_STREAM

    ssl_version = ssl.PROTOCOL_SSLv3
    ssl_cert_reqs = ssl.CERT_REQUIRED

    def __init__(self, address, port, ssl_ca_certs='cacert.pem'):
        self.address = address
        self.port = port
        self._sock = socket.socket(self.socket_family, self.socket_type)
        self.ssl_ca_certs = ssl_ca_certs

    def __enter__(self):
        self._sock.connect( (self.address, self.port) )
        self._sslsock = ssl.wrap_socket(self._sock,
            ssl_version=self.ssl_version, cert_reqs=self.ssl_cert_reqs,
            ca_certs=self.ssl_ca_certs)
        return self._sslsock

    def __exit__(self, type, value, traceback):
        self._sslsock.close()   # Closes plain socket, too


if __name__ == '__main__':
    addr = socket.gethostbyname('www.google.com')
    port = 443

    with EasySSLWrap( addr, port ) as sslsock:
        cert_dict = sslsock.getpeercert()
        print cert_dict