1"""PyOpenSSL utilities including HTTPSSocket class which wraps PyOpenSSL
2SSL connection into a httplib-like interface suitable for use with urllib2
3
4"""
5__author__ = "P J Kershaw"
6__date__ = "21/12/10"
7__copyright__ = "(C) 2012 Science and Technology Facilities Council"
8__license__ = "BSD - see LICENSE file in top-level directory"
9__contact__ = "Philip.Kershaw@stfc.ac.uk"
10__revision__ = '$Id$'
11
12from datetime import datetime
13import logging
14import socket
15from io import BytesIO
16
17from OpenSSL import SSL
18
19log = logging.getLogger(__name__)
20
21
22class SSLSocket(object):
23    """SSL Socket class wraps pyOpenSSL's SSL.Connection class implementing
24    the makefile method so that it is compatible with the standard socket
25    interface and usable with httplib.
26
27    @cvar default_buf_size: default buffer size for recv operations in the
28    makefile method
29    @type default_buf_size: int
30    """
31    default_buf_size = 8192
32
33    def __init__(self, ctx, sock=None):
34        """Create SSL socket object
35
36        @param ctx: SSL context
37        @type ctx: OpenSSL.SSL.Context
38        @param sock: underlying socket object
39        @type sock: socket.socket
40        """
41        if sock is not None:
42            self.socket = sock
43        else:
44            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
45            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
46
47        self.__ssl_conn = SSL.Connection(ctx, self.socket)
48        self.buf_size = self.__class__.default_buf_size
49        self._makefile_refs = 0
50
51    def __del__(self):
52        """Close underlying socket when this object goes out of scope
53        """
54        self.close()
55
56    @property
57    def buf_size(self):
58        """Buffer size for makefile method recv() operations"""
59        return self.__buf_size
60
61    @buf_size.setter
62    def buf_size(self, value):
63        """Buffer size for makefile method recv() operations"""
64        if not isinstance(value, int):
65            raise TypeError('Expecting int type for "buf_size"; '
66                            'got %r instead' % type(value))
67        self.__buf_size = value
68
69    def close(self):
70        """Shutdown the SSL connection and call the close method of the
71        underlying socket"""
72        if self._makefile_refs < 1:
73            try:
74                self.__ssl_conn.shutdown()
75            except (SSL.Error, SSL.SysCallError):
76                # Make errors on shutdown non-fatal
77                pass
78        else:
79            self._makefile_refs -= 1
80
81    def set_shutdown(self, mode):
82        """Set the shutdown state of the Connection.
83        @param mode: bit vector of either or both of SENT_SHUTDOWN and
84        RECEIVED_SHUTDOWN
85        """
86        self.__ssl_conn.set_shutdown(mode)
87
88    def get_shutdown(self):
89        """Get the shutdown state of the Connection.
90        @return: bit vector of either or both of SENT_SHUTDOWN and
91        RECEIVED_SHUTDOWN
92        """
93        return self.__ssl_conn.get_shutdown()
94
95    def bind(self, addr):
96        """bind to the given address - calls method of the underlying socket
97        @param addr: address/port number tuple
98        @type addr: tuple"""
99        self.__ssl_conn.bind(addr)
100
101    def listen(self, backlog):
102        """Listen for connections made to the socket.
103
104        @param backlog: specifies the maximum number of queued connections and
105        should be at least 1; the maximum value is system-dependent (usually 5).
106        @param backlog: int
107        """
108        self.__ssl_conn.listen(backlog)
109
110    def set_accept_state(self):
111        """Set the connection to work in server mode. The handshake will be
112        handled automatically by read/write"""
113        self.__ssl_conn.set_accept_state()
114
115    def accept(self):
116        """Accept an SSL connection.
117
118        @return: pair (ssl, addr) where ssl is a new SSL connection object and
119        addr is the address bound to the other end of the SSL connection.
120        @rtype: tuple
121        """
122        return self.__ssl_conn.accept()
123
124    def set_connect_state(self):
125        """Set the connection to work in client mode. The handshake will be
126        handled automatically by read/write"""
127        self.__ssl_conn.set_connect_state()
128
129    def connect(self, addr):
130        """Call the connect method of the underlying socket and set up SSL on
131        the socket, using the Context object supplied to this Connection object
132        at creation.
133
134        @param addr: address/port number pair
135        @type addr: tuple
136        """
137        self.__ssl_conn.connect(addr)
138
139    def shutdown(self, how):
140        """Send the shutdown message to the Connection.
141
142        @param how: for socket.socket this flag determines whether read, write
143        or both type operations are supported.  OpenSSL.SSL.Connection doesn't
144        support this so this parameter is IGNORED
145        @return: true if the shutdown message exchange is completed and false
146        otherwise (in which case you call recv() or send() when the connection
147        becomes readable/writeable.
148        @rtype: bool
149        """
150        return self.__ssl_conn.shutdown()
151
152    def renegotiate(self):
153        """Renegotiate this connection's SSL parameters."""
154        return self.__ssl_conn.renegotiate()
155
156    def pending(self):
157        """@return: numbers of bytes that can be safely read from the SSL
158        buffer.
159        @rtype: int
160        """
161        return self.__ssl_conn.pending()
162
163    def send(self, data, *flags_arg):
164        """Send data to the socket. Nb. The optional flags argument is ignored.
165        - retained for compatibility with socket.socket interface
166
167        @param data: data to send down the socket
168        @type data: string
169        """
170        return self.__ssl_conn.send(data)
171
172    def sendall(self, data):
173        self.__ssl_conn.sendall(data)
174
175    def recv(self, size=default_buf_size):
176        """Receive data from the Connection.
177
178        @param size: The maximum amount of data to be received at once
179        @type size: int
180        @return: data received.
181        @rtype: string
182        """
183        return self.__ssl_conn.recv(size)
184
185    def setblocking(self, mode):
186        """Set this connection's underlying socket blocking _mode_.
187
188        @param mode: blocking mode
189        @type mode: int
190        """
191        self.__ssl_conn.setblocking(mode)
192
193    def fileno(self):
194        """
195        @return: file descriptor number for the underlying socket
196        @rtype: int
197        """
198        return self.__ssl_conn.fileno()
199
200    def getsockopt(self, *args):
201        """See socket.socket.getsockopt
202        """
203        return self.__ssl_conn.getsockopt(*args)
204
205    def setsockopt(self, *args):
206        """See socket.socket.setsockopt
207
208        @return: value of the given socket option
209        @rtype: int/string
210        """
211        return self.__ssl_conn.setsockopt(*args)
212
213    def state_string(self):
214        """Return the SSL state of this connection."""
215        return self.__ssl_conn.state_string()
216
217    def makefile(self, *args):
218        """Specific to Python socket API and required by httplib: convert
219        response into a file-like object.  This implementation reads using recv
220        and copies the output into a StringIO buffer to simulate a file object
221        for consumption by httplib
222
223        Nb. Ignoring optional file open mode (StringIO is generic and will
224        open for read and write unless a string is passed to the constructor)
225        and buffer size - httplib set a zero buffer size which results in recv
226        reading nothing
227
228        @return: file object for data returned from socket
229        @rtype: cStringIO.StringO
230        """
231        self._makefile_refs += 1
232
233        # Optimisation
234        _buf_size = self.buf_size
235
236        i=0
237        stream = BytesIO()
238        startTime = datetime.utcnow()
239        try:
240            dat = self.__ssl_conn.recv(_buf_size)
241            while dat:
242                i+=1
243                stream.write(dat)
244                dat = self.__ssl_conn.recv(_buf_size)
245
246        except (SSL.ZeroReturnError, SSL.SysCallError):
247            # Connection is closed - assuming here that all is well and full
248            # response has been received.  httplib will catch an error in
249            # incomplete content since it checks the content-length header
250            # against the actual length of data received
251            pass
252
253        if log.getEffectiveLevel() <= logging.DEBUG:
254            log.debug("Socket.makefile %d recv calls completed in %s", i,
255                      datetime.utcnow() - startTime)
256
257        # Make sure to rewind the buffer otherwise consumers of the content will
258        # read from the end of the buffer
259        stream.seek(0)
260
261        return stream
262
263    def getsockname(self):
264        """
265        @return: the socket's own address
266        @rtype:
267        """
268        return self.__ssl_conn.getsockname()
269
270    def getpeername(self):
271        """
272        @return: remote address to which the socket is connected
273        """
274        return self.__ssl_conn.getpeername()
275
276    def get_context(self):
277        '''Retrieve the Context object associated with this Connection. '''
278        return self.__ssl_conn.get_context()
279
280    def get_peer_certificate(self):
281        '''Retrieve the other side's certificate (if any)  '''
282        return self.__ssl_conn.get_peer_certificate()
283