Merge pull request #18 from DK777/tproxysocket

add TProxySocket support

README.md

@@ -70,6 +70,7 @@ Options:
                             Java Thrift body, such as 'request:MyRequest(person:Person(name:joe, id:2))'.
                             Path to a file containing any of the above formats.
 - **-z --zookeeper**       Treat the server address as a Zookeeper instance, and make the request to the service being provided at the given path.
+- **-p --proxy [PROXY]**    Access the service via a proxy (for auth reasons) "proxy host:proxy port"
 - **-c --cleanup**         Delete generated code from filesystem after execution
 - **-j --json**            Print result in JSON format
 - **-i --client_id [client_id]**

setup.py

@@ -16,14 +16,14 @@
 
 config = {
     'name': 'thriftcli',
-    'version': 1.1,
+    'version': 1.2,
     'description': 'Thrift CLI',
     'author': 'Neel Virdy',
     'packages': ['thriftcli'],
     'entry_points': {
         'console_scripts': ['thriftcli = thriftcli.thrift_cli:main']
     },
-    'install_requires': ['kazoo', 'mock', 'thrift', 'twitter.common.rpc', 'coverage'],
+    'install_requires': ['kazoo', 'mock', 'requests_kerberos', 'thrift', 'twitter.common.rpc', 'coverage'],
 }
 
 setup(**config)

tests/data/constants.py

@@ -465,20 +465,21 @@
               (TEST_THRIFT_SERVICE_NAME, TEST_SERVER_HOSTNAME2, TEST_SERVER_PORT2),
               None)
 TEST_CLIENT_ID = "client_abc"
-TEST_CLI_ARGS = [TEST_CLI_NAME, TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH]
+TEST_PROXY = "proxy.abc.com:12345"
+TEST_CLI_ARGS = [TEST_CLI_NAME, TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH, '-p', TEST_PROXY]
 TEST_CLI_ARGS2 = [TEST_CLI_NAME, TEST_ZOOKEEPER_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
                   '-b', TEST_JSON_PATH, '-z', '-j', '-c',  '-i', TEST_CLIENT_ID]
 TEST_CLI_ARGS3 = [TEST_CLI_NAME, TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
-                  '-I', TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2]
+                  '-I', TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2, '--proxy', TEST_PROXY]
 TEST_CLI_ARGS4 = [TEST_CLI_NAME, TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
                   '--body', TEST_JSON_REQUEST_BODY, '--include', TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2,
                   '--client_id', TEST_CLIENT_ID]
 TEST_CLI_ARGS5 = [TEST_CLI_NAME, TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
                   '--body', TEST_INVALID_REQUEST_BODY]
-TEST_PARSED_ARGS = (TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH, [], {}, False, False, False, None)
+TEST_PARSED_ARGS = (TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH, [], {}, False, False, False, None, TEST_PROXY)
 TEST_PARSED_ARGS2 = (TEST_ZOOKEEPER_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH, [],
-                     TEST_ARGUMENT_DICTIONARY, True, True, True, TEST_CLIENT_ID)
+                     TEST_ARGUMENT_DICTIONARY, True, True, True, TEST_CLIENT_ID, None)
 TEST_PARSED_ARGS3 = (TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
-                     [TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2], {}, False, False, False, None)
+                     [TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2], {}, False, False, False, None, TEST_PROXY)
 TEST_PARSED_ARGS4 = (TEST_SERVER_ADDRESS, TEST_THRIFT_ENDPOINT_NAME, TEST_THRIFT_PATH,
-                     [TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2], TEST_ARGUMENT_DICTIONARY, False, False, False, TEST_CLIENT_ID)
+                     [TEST_THRIFT_DIR_PATH, TEST_THRIFT_DIR_PATH2], TEST_ARGUMENT_DICTIONARY, False, False, False, TEST_CLIENT_ID, None)

thriftcli/thrift_cli.py

@@ -33,7 +33,7 @@ class ThriftCLI(object):
 
     """
 
-    def __init__(self, thrift_path, server_address, service_name, thrift_dir_paths=None, zookeeper=False, client_id=None):
+    def __init__(self, thrift_path, server_address, service_name, thrift_dir_paths=None, zookeeper=False, client_id=None, proxy=None):
         """
         :param thrift_path: the path to the thrift file being used.
         :type thrift_path: str
@@ -45,7 +45,8 @@ def __init__(self, thrift_path, server_address, service_name, thrift_dir_paths=N
         :type zookeeper: bool
         :param client_id: Finagle client id for identifying requests
         :type client_id: str
-
+        :param proxy: [<proxy host>:<proxy port>] to route request through
+        :type proxy: str
         """
         self._thrift_path = _find_path(thrift_path)
         self._thrift_argument_converter = ThriftArgumentConverter(self._thrift_path, thrift_dir_paths)
@@ -54,7 +55,7 @@ def __init__(self, thrift_path, server_address, service_name, thrift_dir_paths=N
             server_address = get_server_address(server_address, service_name)
         self._thrift_executor = ThriftExecutor(self._thrift_path, server_address, self._service_reference,
                                                self._thrift_argument_converter._parse_result.namespaces,
-                                               thrift_dir_paths=thrift_dir_paths, client_id=client_id)
+                                               thrift_dir_paths=thrift_dir_paths, client_id=client_id, proxy=proxy)
 
     def run(self, method_name, request_body, return_json=False):
         """ Runs the endpoint on the connected server as defined by the thrift file.
@@ -179,8 +180,9 @@ def _parse_namespace(args):
     return_json = args.json
     cleanup = args.cleanup
     client_id = args.client_id
+    proxy = args.proxy
     return (server_address, endpoint, thrift_path, thrift_dir_paths, request_body, zookeeper, return_json, cleanup,
-            client_id)
+            client_id, proxy)
 
 
 def _make_parser():
@@ -203,6 +205,8 @@ def _make_parser():
                         help='json string or path to json file encoding the request body')
     parser.add_argument('-z', '--zookeeper', action='store_true',
                         help='treat server address as a zookeeper host with a path')
+    parser.add_argument('-p', '--proxy', type=str,
+                        help='access the service via a proxy (for auth reasons) [<proxy host>:<proxy port>]')
     parser.add_argument('-c', '--cleanup', action='store_true',
                         help='remove generated code after execution')
     parser.add_argument('-j', '--json', action='store_true',
@@ -215,7 +219,7 @@ def _make_parser():
 
 
 def _run_cli(server_address, endpoint_name, thrift_path, thrift_dir_paths, request_body, zookeeper, return_json,
-             remove_generated_src, client_id):
+             remove_generated_src, client_id, proxy):
     """ Runs a remote request and prints the result if it is not None.
 
     :param server_address: the address of the Thrift server to request
@@ -234,6 +238,8 @@ def _run_cli(server_address, endpoint_name, thrift_path, thrift_dir_paths, reque
     :type remove_generated_src: bool
     :param client_id: Finagle client id for identifying requests
     :type client_id: str
+    :param proxy: [<proxy host>:<proxy port>] to route request through
+    :type proxy: str
     :param verbose: log details
     :type verbose: bool
 
@@ -248,7 +254,8 @@ def _run_cli(server_address, endpoint_name, thrift_path, thrift_dir_paths, reque
         service_name,
         thrift_dir_paths + environment_defined_paths,
         zookeeper,
-        client_id=client_id
+        client_id=client_id,
+        proxy=proxy
     )
     try:
         result = cli.run(method_name, request_body, return_json)

thriftcli/thrift_executor.py

@@ -22,21 +22,22 @@
 from twitter.common.rpc.finagle.protocol import TFinagleProtocol
 
 from .thrift_cli_error import ThriftCLIError
+from .transport import TProxySocket
 
 
 class ThriftExecutor(object):
     """ This class handles connecting to and communicating with the Thrift server. """
 
     def __init__(self, thrift_path, server_address, service_reference, basename_to_namespaces, thrift_dir_paths=None,
-                 client_id=None):
+                 client_id=None, proxy=None):
         """ Opens a connection with the server and generates then imports the thrift-defined python code.
 
         :param thrift_path: the path to the Thrift file defining the service being requested
         :param server_address: the address to the server implementing the service
         :param service_reference: the namespaced service name in the format <file-name>.<service-name>
         :param thrift_dir_paths: a list of paths to directories containing Thrift file dependencies
         :param client_id: Finagle client id for identifying requests
-
+        :param proxy: [<proxy host>:<proxy port>] to route request through
         """
         self._thrift_path = thrift_path
         self._server_address = server_address
@@ -48,6 +49,7 @@ def __init__(self, thrift_path, server_address, service_reference, basename_to_n
 
         self._client_id = client_id
         self._service_reference = service_reference
+        self._proxy = proxy
         self._open_connection(server_address)
         self._generate_and_import_packages(basename_to_namespaces)
 
@@ -128,7 +130,11 @@ def _open_connection(self, address):
 
         """
         (url, port) = self._parse_address_for_hostname_and_port(address)
-        self._transport = TSocket.TSocket(url, port)
+        if self._proxy:
+            proxy_host, proxy_port = self._proxy.split(":")
+            self._transport = TProxySocket(proxy_host, proxy_port, url, port)
+        else:
+            self._transport = TSocket.TSocket(url, port)
         self._transport = TTransport.TFramedTransport(self._transport)
         self._transport.open()
         self._protocol = TFinagleProtocol(self._transport, client_id=self._client_id)

thriftcli/transport.py

@@ -0,0 +1,66 @@
+"""For those times you want to access your TSocket transport service via a proxy.
+
+Maybe you need to access some backend services via a squid proxy that uses kerberos for authentication, authorization
+and auditing ofthese requests.
+
+TProxySocket is a thin wrapper around TSocket transport that uses httplib to handle setting up the tunnel 
+and requests_kerberos to handle kerberos auth.
+"""
+import httplib
+import requests_kerberos
+import socket
+
+from thrift.transport import TSocket
+from thrift.transport import TTransport
+
+class TProxySocket(TSocket.TSocket):
+  """Thrift transport, adds proxy support to TSocket transport."""
+  def __init__(self, proxy_host=None, proxy_port=None, *args, **kwargs):
+    if proxy_host is None and proxy_port is None:
+      return TSocket.TSocket(*args, **kwargs)
+
+    TSocket.TSocket.__init__(self, *args, **kwargs)
+    self.proxy_host = proxy_host
+    self.proxy_port = proxy_port
+
+  def open(self):
+    """Open a connection.
+
+    This is mostly copy pasta from thrift.transport.TSocket.open, but adds call to _setup_tunnel().
+    """
+    try:
+      res0 = self._resolveAddr()
+      for res in res0:
+        try:
+          self.handle = self._setup_tunnel(res[4])
+          self.handle.settimeout(self._timeout)
+        except socket.error as e:
+          if res is not res0[-1]:
+            continue
+          else:
+            raise e
+        break
+    except socket.error as e:
+      if self._unix_socket:
+        message = 'Could not connect to socket %s' % self._unix_socket
+      else:
+        message = 'Could not connect to %s:%d' % (self.host, self.port)
+      raise TTransport.TTransportException(type=TTransport.TTransportException.NOT_OPEN,
+                                message=message)
+
+  def _setup_tunnel(self, host_port):
+    """Use HTTPConnection to HTTP CONNECT to our proxy, connect to the backend and return socket for this tunnel.
+
+    host_port: tuple (host, port) from thrift.transport.TSocket._resolveAddr.
+    """
+    conn = httplib.HTTPConnection(self.proxy_host, self.proxy_port)
+    auth_header = requests_kerberos.HTTPKerberosAuth().generate_request_header(None,
+                                                                                self.proxy_host,
+                                                                                is_preemptive=True)
+    conn.set_tunnel(
+                    *host_port,
+                    headers={'Proxy-Authorization': auth_header})
+    conn.connect()
+    return conn.sock
+
+