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