Modify return object for retrieve_key(key_id, twsk)

Modify the return type of the function retrieve_key(key_id,
trans_wrapped_session_key) from returining a tuple KeyData, unwrapped_key
to KeyData by setting the unwrapped_key to KeyData.private_data attribute
for the case where trans_wrapped_session_key is not provided by the caller.

2 files changed, 27 insertions, 18 deletions

diff --git a/base/common/python/pki/key.py b/base/common/python/pki/key.py
index d9113cd3f..69e01c852 100644
--- a/base/common/python/pki/key.py
+++ b/base/common/python/pki/key.py
@@ -63,6 +63,10 @@ class KeyData(object):
         self.nonceData = None
         self.size = None
         self.wrappedPrivateData = None
+        
+        # To store the unwrapped key information.
+        # Is not transferred in the response.
+        self.private_data = None
 
     @classmethod
     def from_json(cls, attr_list):
@@ -427,7 +431,7 @@ class KeyClient(object):
             raise TypeError("Request ID must be specified")
 
         url = self.key_requests_url + '/' + request_id + '/approve'
-        self.connection.post(url, self.headers)
+        self.connection.post(url, None, self.headers)
 
     @pki.handle_exceptions()
     def reject_request(self, request_id):
@@ -436,7 +440,7 @@ class KeyClient(object):
             raise TypeError("Request ID must be specified")
 
         url = self.key_requests_url + '/' + request_id + '/reject'
-        self.connection.post(url, self.headers)
+        self.connection.post(url, None, self.headers)
 
     @pki.handle_exceptions()
     def cancel_request(self, request_id):
@@ -445,7 +449,7 @@ class KeyClient(object):
             raise TypeError("Request ID must be specified")
 
         url = self.key_requests_url + '/' + request_id + '/cancel'
-        self.connection.post(url, self.headers)
+        self.connection.post(url, None, self.headers)
 
     @pki.handle_exceptions()
     def create_request(self, request):
@@ -687,14 +691,15 @@ class KeyClient(object):
         1) trans_wrapped_session_key is not provided by caller.
 
         In this case, the function will call CryptoUtil methods to generate and wrap the
-        session key.  The function will return the tuple (KeyData, unwrapped_secret)
+        session key.  The function will return the KeyData object with a private_data attribute
+        which stores the unwrapped key information.
 
         2)  The trans_wrapped_session_key is provided by the caller.
 
         In this case, the function will simply pass the data to the DRM, and will return the secret
         wrapped in the session key.  The secret will still need to be unwrapped by the caller.
 
-        The function will return the tuple (KeyData, None), where the KeyData structure includes the
+        The function will return the KeyData object, where the KeyData structure includes the
         wrapped secret and some nonce data to be used as a salt when unwrapping.
         '''
         if key_id is None:
@@ -717,14 +722,12 @@ class KeyClient(object):
                         trans_wrapped_session_key=base64.encodestring(trans_wrapped_session_key))
 
         key_data = self.retrieve_key_data(request)
-        if key_provided:
-            return key_data, None
-
-        unwrapped_key = self.crypto.symmetric_unwrap(
+        if not key_provided:
+            key_data.private_data = self.crypto.symmetric_unwrap(
                                 key_data.wrappedPrivateData,
                                 session_key,
                                 nonce_iv=key_data.nonceData)
-        return key_data, unwrapped_key
+        return key_data
 
     @pki.handle_exceptions()
     def retrieve_key_by_passphrase(self, key_id, passphrase=None,
diff --git a/base/kra/functional/drmtest.py b/base/kra/functional/drmtest.py
index 00f7d7a9d..5b09b0aaa 100644
--- a/base/kra/functional/drmtest.py
+++ b/base/kra/functional/drmtest.py
@@ -60,6 +60,8 @@ def print_key_data(key_data):
     print "Key Size: " + str(key_data.size)
     print "Nonce Data: " + base64.encodestring(key_data.nonceData)
     print "Wrapped Private Data: " + base64.encodestring(key_data.wrappedPrivateData)
+    if key_data.private_data is not None:
+        print "Private Data: " + base64.encodestring(key_data.private_data)
 
 def main():
     ''' test code execution '''
@@ -93,8 +95,12 @@ def main():
 
     # Test 2: Get key request info
     print "Now getting key request"
-    keyrequest = keyclient.get_request_info('2')
-    print_key_request(keyrequest)
+    try:
+        keyrequest = keyclient.get_request_info('2')
+        print_key_request(keyrequest)
+    except pki.RequestNotFoundException as e:
+        print "Can be ignored for a first time run!!"
+        print e
 
     # Test 3: List requests
     print "Now listing some requests"
@@ -132,7 +138,7 @@ def main():
     session_key = crypto.generate_session_key()
     wrapped_session_key = crypto.asymmetric_wrap(session_key, keyclient.transport_cert)
     print "My key id is " + str(key_id)
-    key_data, _unwrapped_key = keyclient.retrieve_key(key_id, trans_wrapped_session_key=wrapped_session_key)
+    key_data = keyclient.retrieve_key(key_id, trans_wrapped_session_key=wrapped_session_key)
     print_key_data(key_data)
     unwrapped_key = crypto.symmetric_unwrap(key_data.wrappedPrivateData,
                                             session_key,
@@ -140,9 +146,9 @@ def main():
     key1 = base64.encodestring(unwrapped_key)
 
     # Test 7: Recover key without providing trans_wrapped_session_key
-    key_data, unwrapped_key = keyclient.retrieve_key(key_id)
+    key_data = keyclient.retrieve_key(key_id)
     print_key_data(key_data)
-    key2 = base64.encodestring(unwrapped_key)
+    key2 = base64.encodestring(key_data.private_data)
 
     # Test 8 - Confirm that keys returned are the same
     if key1 == key2:
@@ -172,7 +178,7 @@ def main():
     # Test 12 - Test exception on retrieve_key.
     print "Try to retrieve an invalid key"
     try:
-        key_data, unwrapped_key = keyclient.retrieve_key('2000003434')
+        key_data = keyclient.retrieve_key('2000003434')
     except pki.KeyNotFoundException as exc:
         print "KeyNotFoundException thrown - Code:" + exc.code + " Message: " + exc.message
 
@@ -226,9 +232,9 @@ def main():
     key_info = keyclient.get_active_key_info(client_key_id)
     print_key_info(key_info)
 
-    key_data, unwrapped_key = keyclient.retrieve_key(key_info.get_key_id())
+    key_data = keyclient.retrieve_key(key_info.get_key_id())
     print_key_data(key_data)
-    key2 = base64.encodestring(unwrapped_key)
+    key2 = base64.encodestring(key_data.private_data)
 
     if key1 == key2:
         print "Success: archived and recovered keys match"