>From 72327ded1a27ca8fc35755fe34fc68d249f0b5d5 Mon Sep 17 00:00:00 2001
From: Abhishek Koneru <akoneru@redhat.com>
Date: Mon, 3 Mar 2014 07:54:04 -0500
Subject: [PATCH] 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.
---
 base/common/python/pki/key.py  | 13 +++++++++----
 base/kra/functional/drmtest.py |  8 +++++---
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/base/common/python/pki/key.py b/base/common/python/pki/key.py
index cdcc7380fce3f6b92a3ba5c7e8c7aecf1ac4fe9d..2adb5dfd6eaca5a4aa0b95d5453fe12be023531b 100644
--- a/base/common/python/pki/key.py
+++ b/base/common/python/pki/key.py
@@ -63,6 +63,9 @@ 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):
@@ -693,14 +696,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:
@@ -724,13 +728,14 @@ class KeyClient(object):
 
         key_data = self.retrieve_key_data(request)
         if key_provided:
-            return key_data, None
+            return key_data
 
         unwrapped_key = self.crypto.symmetric_unwrap(
                                 key_data.wrappedPrivateData,
                                 session_key,
                                 nonce_iv=key_data.nonceData)
-        return key_data, unwrapped_key
+        key_data.private_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 169a7879fab4cda61fc915d3ca2652039c3c5929..2d4f7565638d23bbf4c1da1b0948235c6c360e23 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 '''
@@ -132,7 +134,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 +142,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:
-- 
1.8.5.3

