diff options
Diffstat (limited to 'pyload/remote/thriftbackend/ThriftClient.py')
-rw-r--r-- | pyload/remote/thriftbackend/ThriftClient.py | 87 |
1 files changed, 87 insertions, 0 deletions
diff --git a/pyload/remote/thriftbackend/ThriftClient.py b/pyload/remote/thriftbackend/ThriftClient.py new file mode 100644 index 000000000..d7b0d2919 --- /dev/null +++ b/pyload/remote/thriftbackend/ThriftClient.py @@ -0,0 +1,87 @@ +# -*- coding: utf-8 -*- + +import sys +from socket import error +from os.path import dirname, abspath, join +from traceback import print_exc + +try: + import thrift +except ImportError: + sys.path.append(abspath(join(dirname(abspath(__file__)), "..", "..", "lib"))) + +from thrift.transport import TTransport +#from thrift.transport.TZlibTransport import TZlibTransport +from Socket import Socket +from Protocol import Protocol + +# modules should import ttypes from here, when want to avoid importing API + +from thriftgen.pyload import Pyload +from thriftgen.pyload.ttypes import * + +ConnectionClosed = TTransport.TTransportException + +class WrongLogin(Exception): + pass + +class NoConnection(Exception): + pass + +class NoSSL(Exception): + pass + +class ThriftClient(object): + def __init__(self, host="localhost", port=7227, user="", password=""): + + self.createConnection(host, port) + try: + self.transport.open() + except error, e: + if e.args and e.args[0] in (111, 10061): + raise NoConnection + else: + print_exc() + raise NoConnection + + try: + correct = self.client.login(user, password) + except error, e: + if e.args and e.args[0] == 104: + #connection reset by peer, probably wants ssl + try: + self.createConnection(host, port, True) + #set timeout or a ssl socket will block when querying none ssl server + self.socket.setTimeout(10) + + except ImportError: + #@TODO untested + raise NoSSL + try: + self.transport.open() + correct = self.client.login(user, password) + finally: + self.socket.setTimeout(None) + elif e.args and e.args[0] == 32: + raise NoConnection + else: + print_exc() + raise NoConnection + + if not correct: + self.transport.close() + raise WrongLogin + + def createConnection(self, host, port, ssl=False): + self.socket = Socket(host, port, ssl) + self.transport = TTransport.TBufferedTransport(self.socket) +# self.transport = TZlibTransport(TTransport.TBufferedTransport(self.socket)) + + protocol = Protocol(self.transport) + self.client = Pyload.Client(protocol) + + def close(self): + self.transport.close() + + def __getattr__(self, item): + return getattr(self.client, item) |