Added "headers" argument to "getClient()" methods to add arbitrary request headers ZTK-1.1
authorThierry Florac <tflorac@ulthar.net>
Thu, 30 May 2013 18:35:15 +0200
branchZTK-1.1
changeset 223 85a03fd95686
parent 222 ac6fb60952c9
child 224 672aebb49077
Added "headers" argument to "getClient()" methods to add arbitrary request headers
src/ztfy/utils/protocol/xmlrpc.py
--- a/src/ztfy/utils/protocol/xmlrpc.py	Fri May 24 11:41:36 2013 +0200
+++ b/src/ztfy/utils/protocol/xmlrpc.py	Thu May 30 18:35:15 2013 +0200
@@ -53,12 +53,13 @@
     _http_connection = httplib.HTTPConnection
     _http_connection_compat = TimeoutHTTP
 
-    def __init__(self, user_agent, credentials=(), cookies=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
+    def __init__(self, user_agent, credentials=(), cookies=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, headers=None):
         xmlrpclib.Transport.__init__(self)
         self.user_agent = user_agent
         self.credentials = credentials
         self.cookies = cookies
         self.timeout = timeout
+        self.headers = headers
         if self._connection_required_compat():
             self.make_connection = self._make_connection_compat
             self.get_response = self._get_response_compat
@@ -103,6 +104,11 @@
             auth = 'Basic %s' % base64.encodestring("%s:%s" % self.credentials).strip()
             connection.putheader('Authorization', auth)
 
+    # send custom headers
+    def send_headers(self, connection):
+        for k, v in (self.headers or {}).iteritems():
+            connection.putheader(k, v)
+
     # dummy request class for extracting cookies
     class CookieRequest(urllib2.Request):
         pass
@@ -137,6 +143,7 @@
         self.send_request(connection, handler, request_body)
         self.send_host(connection, host)
         self.send_user_agent(connection)
+        self.send_headers(connection)
         self.send_content(connection, request_body)
         # get response
         return self.get_response(connection, host, handler)
@@ -175,19 +182,19 @@
     _http_connection_compat = TimeoutHTTPS
 
 
-def getClient(uri, credentials=(), verbose=False, allow_none=0, timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
+def getClient(uri, credentials=(), verbose=False, allow_none=0, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, headers=None):
     """Get an XML-RPC client which supports basic authentication"""
     if uri.startswith('https:'):
-        transport = SecureXMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY secure transport)', credentials, timeout=timeout)
+        transport = SecureXMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY secure transport)', credentials, timeout=timeout, headers=headers)
     else:
-        transport = XMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY basic transport)', credentials, timeout=timeout)
+        transport = XMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY basic transport)', credentials, timeout=timeout, headers=headers)
     return xmlrpclib.Server(uri, transport=transport, verbose=verbose, allow_none=allow_none)
 
 
-def getClientWithCookies(uri, credentials=(), verbose=False, allow_none=0, timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
-    """Get an XML-RPC client which supports authentication throught cookies"""
+def getClientWithCookies(uri, credentials=(), verbose=False, allow_none=0, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, headers=None):
+    """Get an XML-RPC client which supports authentication through cookies"""
     if uri.startswith('https:'):
-        transport = SecureXMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY secure cookie transport)', credentials, cookielib.CookieJar(), timeout)
+        transport = SecureXMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY secure cookie transport)', credentials, cookielib.CookieJar(), timeout, headers)
     else:
-        transport = XMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY secure cookie transport)', credentials, cookielib.CookieJar(), timeout)
+        transport = XMLRPCCookieAuthTransport('Python XML-RPC Client/0.1 (ZTFY basic cookie transport)', credentials, cookielib.CookieJar(), timeout, headers)
     return xmlrpclib.Server(uri, transport=transport, verbose=verbose, allow_none=allow_none)