1# -*- coding: utf-8 -*-
2
3from __future__ import absolute_import
4
5import contextlib
6import socket
7import warnings
8
9from thriftpy2._compat import PY3, PY35
10if PY3:
11    import urllib
12else:
13    import urllib2 as urllib
14    import urlparse
15    urllib.parse = urlparse
16    urllib.parse.quote = urllib.quote
17
18from thriftpy2.protocol import TBinaryProtocolFactory
19from thriftpy2.server import TThreadedServer
20from thriftpy2.thrift import TProcessor, TClient
21from thriftpy2.transport import (
22    TBufferedTransportFactory,
23    TServerSocket,
24    TSSLServerSocket,
25    TSocket,
26    TSSLSocket,
27)
28
29
30def make_client(service, host="localhost", port=9090, unix_socket=None,
31                proto_factory=TBinaryProtocolFactory(),
32                trans_factory=TBufferedTransportFactory(),
33                timeout=3000, cafile=None, ssl_context=None, certfile=None,
34                keyfile=None, url="", socket_family=socket.AF_INET):
35    if url:
36        parsed_url = urllib.parse.urlparse(url)
37        host = parsed_url.hostname or host
38        port = parsed_url.port or port
39    if unix_socket:
40        socket = TSocket(unix_socket=unix_socket, socket_timeout=timeout)
41        if certfile:
42            warnings.warn("SSL only works with host:port, not unix_socket.")
43    elif host and port:
44        if cafile or ssl_context:
45            socket = TSSLSocket(host, port, socket_timeout=timeout,
46                                socket_family=socket_family, cafile=cafile,
47                                certfile=certfile, keyfile=keyfile,
48                                ssl_context=ssl_context)
49        else:
50            socket = TSocket(host, port, socket_family=socket_family, socket_timeout=timeout)
51    else:
52        raise ValueError("Either host/port or unix_socket or url must be provided.")
53
54    transport = trans_factory.get_transport(socket)
55    protocol = proto_factory.get_protocol(transport)
56    transport.open()
57    return TClient(service, protocol)
58
59
60def make_server(service, handler,
61                host="localhost", port=9090, unix_socket=None,
62                proto_factory=TBinaryProtocolFactory(),
63                trans_factory=TBufferedTransportFactory(),
64                client_timeout=3000, certfile=None):
65    processor = TProcessor(service, handler)
66
67    if unix_socket:
68        server_socket = TServerSocket(unix_socket=unix_socket)
69        if certfile:
70            warnings.warn("SSL only works with host:port, not unix_socket.")
71    elif host and port:
72        if certfile:
73            server_socket = TSSLServerSocket(
74                host=host, port=port, client_timeout=client_timeout,
75                certfile=certfile)
76        else:
77            server_socket = TServerSocket(
78                host=host, port=port, client_timeout=client_timeout)
79    else:
80        raise ValueError("Either host/port or unix_socket must be provided.")
81
82    server = TThreadedServer(processor, server_socket,
83                             iprot_factory=proto_factory,
84                             itrans_factory=trans_factory)
85    return server
86
87
88@contextlib.contextmanager
89def client_context(service, host="localhost", port=9090, unix_socket=None,
90                   proto_factory=TBinaryProtocolFactory(),
91                   trans_factory=TBufferedTransportFactory(),
92                   timeout=None, socket_timeout=3000, connect_timeout=3000,
93                   cafile=None, ssl_context=None, certfile=None, keyfile=None,
94                   url=""):
95    if url:
96        parsed_url = urllib.parse.urlparse(url)
97        host = parsed_url.hostname or host
98        port = parsed_url.port or port
99
100    if timeout:
101        warnings.warn("`timeout` deprecated, use `socket_timeout` and "
102                      "`connect_timeout` instead.")
103        socket_timeout = connect_timeout = timeout
104
105    if unix_socket:
106        socket = TSocket(unix_socket=unix_socket,
107                         connect_timeout=connect_timeout,
108                         socket_timeout=socket_timeout)
109        if certfile:
110            warnings.warn("SSL only works with host:port, not unix_socket.")
111    elif host and port:
112        if cafile or ssl_context:
113            socket = TSSLSocket(host, port,
114                                connect_timeout=connect_timeout,
115                                socket_timeout=socket_timeout,
116                                cafile=cafile,
117                                certfile=certfile, keyfile=keyfile,
118                                ssl_context=ssl_context)
119        else:
120            socket = TSocket(host, port,
121                             connect_timeout=connect_timeout,
122                             socket_timeout=socket_timeout)
123    else:
124        raise ValueError("Either host/port or unix_socket or url must be provided.")
125
126    try:
127        transport = trans_factory.get_transport(socket)
128        protocol = proto_factory.get_protocol(transport)
129        transport.open()
130        yield TClient(service, protocol)
131
132    finally:
133        transport.close()
134
135
136if PY35:
137    from thriftpy2.contrib.aio.rpc import (
138        make_server as make_aio_server,
139        make_client as make_aio_client
140    )
141