blob: 8ab91ca61d2f6e7a80aaba878b5004221726857e [file] [log] [blame]
Bryan Duxbury2b969ad2011-02-22 18:20:53 +00001import sys
2sys.path.append('/usr/lib/python2.6/site-packages/')
3
4from thrift.transport import TSocket
5import socket, ssl
6
7class TSSLSocket(TSocket.TSocket):
8 def open(self):
9 try:
10 res0 = self._resolveAddr()
11 for res in res0:
12 plain_sock = socket.socket(res[0], res[1])
13 #TODO verify server cert
14 self.handle = ssl.wrap_socket(plain_sock, ssl_version=ssl.PROTOCOL_TLSv1)
15 self.handle.settimeout(self._timeout)
16 try:
17 self.handle.connect(res[4])
18 except socket.error, e:
19 if res is not res0[-1]:
20 continue
21 else:
22 raise e
23 break
24 except socket.error, e:
25 if self._unix_socket:
26 message = 'Could not connect to secure socket %s' % self._unix_socket
27 else:
28 message = 'Could not connect to %s:%d' % (self.host, self.port)
29 raise TTransportException(type=TTransportException.NOT_OPEN, message=message)
30
31class TSSLServerSocket(TSocket.TServerSocket):
32 def accept(self):
33 plain_client, addr = self.handle.accept()
34 result = TSocket.TSocket()
35 #TODO take certfile/keyfile as a parameter at setup
36 client = ssl.wrap_socket(plain_client, certfile='cert.pem', server_side=True)
37 result.setHandle(client)
38 return result