Move the cli RPC authorization_headers code into rpc_client_lib.
Make autotest_lib.server.frontend use it for RPC authentication.

Signed-off-by: Gregory Smith <gps@google.com>


git-svn-id: http://test.kernel.org/svn/autotest/trunk@4231 592f7852-d20e-0410-864c-8624ca9c26a4
diff --git a/cli/rpc.py b/cli/rpc.py
index aea2b20..1b50ea7 100644
--- a/cli/rpc.py
+++ b/cli/rpc.py
@@ -2,7 +2,7 @@
 # Copyright 2008 Google Inc. All Rights Reserved.
 #
 
-import os, getpass
+import os
 from autotest_lib.frontend.afe import rpc_client_lib
 from autotest_lib.frontend.afe.json_rpc import proxy
 from autotest_lib.client.common_lib import global_config, utils
@@ -29,20 +29,6 @@
     return web_server
 
 
-def base_authorization_headers(username, server):
-    if not username:
-        if 'AUTOTEST_USER' in os.environ:
-            username = os.environ['AUTOTEST_USER']
-        else:
-            username = getpass.getuser()
-    return {'AUTHORIZATION' : username}
-
-
-authorization_headers = utils.import_site_function(
-        __file__, 'autotest_lib.cli.site_rpc', 'authorization_headers',
-        base_authorization_headers)
-
-
 class rpc_comm(object):
     """Shared AFE/TKO RPC class stuff"""
     def __init__(self, web_server, rpc_path, username):
@@ -54,7 +40,8 @@
     def _connect(self, rpc_path):
         # This does not fail even if the address is wrong.
         # We need to wait for an actual RPC to fail
-        headers = authorization_headers(self.username, self.web_server)
+        headers = rpc_client_lib.authorization_headers(self.username,
+                                                       self.web_server)
         rpc_server = self.web_server + rpc_path
         return rpc_client_lib.get_proxy(rpc_server, headers=headers)
 
diff --git a/frontend/afe/rpc_client_lib.py b/frontend/afe/rpc_client_lib.py
index 1707e04..58268e3 100644
--- a/frontend/afe/rpc_client_lib.py
+++ b/frontend/afe/rpc_client_lib.py
@@ -1,11 +1,35 @@
-"""\
-This module provides a get_proxy() function, which should be used to access
-the afe RPC server.
+"""
+This module provides utility functions useful when writing clients for the RPC
+server.
 """
 
 __author__ = 'showard@google.com (Steve Howard)'
 
+import getpass, os
 from json_rpc import proxy
+from autotest_lib.client.common_lib import utils
+
 
 def get_proxy(*args, **kwargs):
+    """Use this to access the AFE or TKO RPC interfaces."""
     return proxy.ServiceProxy(*args, **kwargs)
+
+
+def _base_authorization_headers(username, server):
+    """
+    Don't call this directly, call authorization_headers().
+    This implementation may be overridden by site code.
+
+    @returns A dictionary of authorization headers to pass in to get_proxy().
+    """
+    if not username:
+        if 'AUTOTEST_USER' in os.environ:
+            username = os.environ['AUTOTEST_USER']
+        else:
+            username = getpass.getuser()
+    return {'AUTHORIZATION' : username}
+
+
+authorization_headers = utils.import_site_function(
+        __file__, 'autotest_lib.frontend.afe.site_rpc_client_lib',
+        'authorization_headers', _base_authorization_headers)
diff --git a/server/frontend.py b/server/frontend.py
index 23bc855..beafd23 100644
--- a/server/frontend.py
+++ b/server/frontend.py
@@ -9,8 +9,8 @@
 We turn the JSON dictionaries into real objects that are more idiomatic
 
 For docs, see:
-    http://autotest/afe/server/noauth/rpc/
-    http://autotest/new_tko/server/noauth/rpc/
+    http://autotest/afe/server/rpc_doc/
+    http://autotest/new_tko/server/rpc_doc/
     http://docs.djangoproject.com/en/dev/ref/models/querysets/#queryset-api
 """
 
@@ -71,8 +71,9 @@
         self.print_log = print_log
         self.debug = debug
         self.reply_debug = reply_debug
-        headers = {'AUTHORIZATION' : self.user}
-        rpc_server = 'http://' + server + path
+        http_server = 'http://' + server
+        headers = rpc_client_lib.authorization_headers(user, http_server)
+        rpc_server = http_server + path
         if debug:
             print 'SERVER: %s' % rpc_server
             print 'HEADERS: %s' % headers
@@ -104,7 +105,7 @@
 class TKO(RpcClient):
     def __init__(self, user=None, server=None, print_log=True, debug=False,
                  reply_debug=False):
-        super(TKO, self).__init__(path='/new_tko/server/noauth/rpc/',
+        super(TKO, self).__init__(path='/new_tko/server/rpc/',
                                   user=user,
                                   server=server,
                                   print_log=print_log,
@@ -123,7 +124,7 @@
     def __init__(self, user=None, server=None, print_log=True, debug=False,
                  reply_debug=False, job=None):
         self.job = job
-        super(AFE, self).__init__(path='/afe/server/noauth/rpc/',
+        super(AFE, self).__init__(path='/afe/server/rpc/',
                                   user=user,
                                   server=server,
                                   print_log=print_log,
diff --git a/server/frontend_unittest.py b/server/frontend_unittest.py
index d5b19dc..e72ccf7 100755
--- a/server/frontend_unittest.py
+++ b/server/frontend_unittest.py
@@ -37,6 +37,9 @@
     def test_init(self):
         os.environ['LOGNAME'] = 'unittest-user'
         GLOBAL_CONFIG.override_config_value('SERVER', 'hostname', 'test-host')
+        rpc_client_lib.authorization_headers.expect_call(
+                'unittest-user', 'http://test-host').and_return(
+                        {'AUTHORIZATION': 'unittest-user'})
         rpc_client_lib.get_proxy.expect_call(
                 'http://test-host/path',
                 headers={'AUTHORIZATION': 'unittest-user'})
@@ -52,8 +55,11 @@
             id = 'idFoo'
             results_platform_map = {'NORAD' : {'Seeking_Joshua': ['WOPR']}}
         GLOBAL_CONFIG.override_config_value('SERVER', 'hostname', 'chess')
+        rpc_client_lib.authorization_headers.expect_call(
+                'david', 'http://chess').and_return(
+                        {'AUTHORIZATION': 'david'})
         rpc_client_lib.get_proxy.expect_call(
-                'http://chess/afe/server/noauth/rpc/',
+                'http://chess/afe/server/rpc/',
                 headers={'AUTHORIZATION': 'david'})
         self.god.stub_function(utils, 'send_email')
         utils.send_email.expect_any_call()