>From 11f46f3d508a86a12fa73658764be90cc25baaf0 Mon Sep 17 00:00:00 2001
From: Ade Lee <alee@redhat.com>
Date: Fri, 21 Feb 2014 11:48:54 -0500
Subject: [PATCH 202/203] Moved key functions out of kraclient.py

---
 base/common/python/pki/cryptoutil.py |   7 +-
 base/common/python/pki/key.py        | 366 +++++++++++++++++++++++++----------
 base/common/python/pki/kraclient.py  | 199 +------------------
 base/kra/functional/drmtest.py       |  55 +++---
 4 files changed, 302 insertions(+), 325 deletions(-)

diff --git a/base/common/python/pki/cryptoutil.py b/base/common/python/pki/cryptoutil.py
index b450e820c29bf9c59332df66a90ea93dac91484c..b5d5fdc1368b51de302ebe2f35110d36cadc22cf 100644
--- a/base/common/python/pki/cryptoutil.py
+++ b/base/common/python/pki/cryptoutil.py
@@ -41,6 +41,11 @@ class CryptoUtil(object):
         pass
 
     @abc.abstractmethod
+    def initialize(self):
+        ''' Initialization code '''
+        pass
+
+    @abc.abstractmethod
     def generate_symmetric_key(self, mechanism=None):
         ''' Generate and return a symmetric key '''
         pass
@@ -112,7 +117,7 @@ class NSSCryptoUtil(CryptoUtil):
         self.certdb_password = certdb_password
         self.nonce_iv = "e4:bb:3b:d3:c3:71:2e:58"
 
-    def initialize_db(self):
+    def initialize(self):
         ''' initialize the nss db.  Must be done before any crypto operations '''
         nss.nss_init(self.certdb_dir)
 
diff --git a/base/common/python/pki/key.py b/base/common/python/pki/key.py
index 0d1dd36f3cc58b7ac3eff15c78bd19a518f5c953..7d93e783ab68f2fc351dbaf4b784b9933673367d 100644
--- a/base/common/python/pki/key.py
+++ b/base/common/python/pki/key.py
@@ -23,6 +23,7 @@
 Module containing the Python client classes for the KeyClient and
 KeyRequestClient REST API on a DRM
 '''
+import base64
 import pki.encoder as encoder
 import json
 import pki
@@ -287,13 +288,24 @@ class KeyClient(object):
     PASS_PHRASE_TYPE = "passPhrase"
     ASYMMETRIC_KEY_TYPE = "asymmetricKey"
 
-    def __init__(self, connection):
+    def __init__(self, connection, crypto, transport_cert_nick=None):
         ''' Constructor '''
         self.connection = connection
         self.headers = {'Content-type': 'application/json',
                         'Accept': 'application/json'}
         self.keyURL = '/rest/agent/keys'
         self.keyRequestsURL = '/rest/agent/keyrequests'
+        self.crypto = crypto
+
+        if transport_cert_nick != None:
+            self.crypto.initialize()
+            self.transport_cert = crypto.get_cert(transport_cert_nick)
+        else:
+            self.transport_cert = None
+
+    def set_transport_cert(self, transport_cert_nick):
+        ''' Set the transport certificate for crypto operations '''
+        self.transport_cert = self.crypto.get_cert(transport_cert_nick)
 
     @pki.handle_exceptions()
     def list_keys(self, client_key_id=None, status=None, max_results=None,
@@ -310,55 +322,6 @@ class KeyClient(object):
         return KeyInfoCollection.from_json(response.json())
 
     @pki.handle_exceptions()
-    def retrieve_key(self, data):
-        ''' Retrieve a secret from the DRM.
-
-            @param: data - a KeyRecoveryRequest containing the keyId of the
-            secret being retrieved, the request_id of the approved recovery
-            request and a wrapping mechanism.  More details at
-            KRAClient.retrieve_key.
-
-            Returns a KeyData object containing the wrapped secret.
-        '''
-        url = self.keyURL + '/retrieve'
-        keyRequest = json.dumps(data, cls=encoder.CustomTypeEncoder, sort_keys=True)
-        response = self.connection.post(url, keyRequest, self.headers)
-        return KeyData.from_json(response.json())
-
-    @pki.handle_exceptions()
-    def request_key_retrieval(self, key_id, request_id, trans_wrapped_session_key=None,
-                     session_wrapped_passphrase=None, passphrase=None, nonce_data=None):
-        ''' Retrieve a secret from the DRM.
-
-            The secret (which is referenced by key_id) can be retrieved only if the
-            recovery request (referenced by request_id) is approved.  key_id and request_id
-            are required.
-
-            Data must be provided to wrap the recovered secret.  This can either be
-            a) a 56-bit DES3 symmetric key, wrapped by the DRM transport key, and
-               passed in trans_wrapped_session_key
-            b) a passphrase.  In this case, the passphrase must be wrapped by a 56-bit
-               symmetric key ("the session key" and passed in session_wrapped_passphrase,
-               and the session key must be wrapped by the DRM transport key and passed
-               in trans_wrapped_session_key
-            c) a passphrase for a p12 file.  If the key being recovered is an asymmetric
-               key, then it is possible to pass in the passphrase for the P12 file to
-               be generated.  This is passed in as passphrase
-
-            nonce_data may also be passed as a salt.
-
-            Returns a KeyData object containing the wrapped secret.
-        '''
-        request = KeyRecoveryRequest(key_id=key_id,
-                                     request_id=request_id,
-                                     trans_wrapped_session_key=trans_wrapped_session_key,
-                                     session_wrapped_passphrase=session_wrapped_passphrase,
-                                     nonce_data=nonce_data,
-                                     passphrase=passphrase)
-
-        return self.retrieve_key(request)
-
-    @pki.handle_exceptions()
     def list_requests(self, request_state=None, request_type=None, client_key_id=None,
                      start=None, page_size=None, max_results=None, max_time=None):
         ''' List/Search key requests in the DRM.
@@ -381,6 +344,72 @@ class KeyClient(object):
         return KeyRequestInfo.from_json(response.json())
 
     @pki.handle_exceptions()
+    def get_key_info(self, key_id):
+        ''' Get the info in the KeyRecord for a specific secret in the DRM. '''
+        url = self.keyURL + '/' + key_id
+        response = self.connection.get(url, headers=self.headers)
+        return KeyInfo.from_json(response.json())
+
+    @pki.handle_exceptions()
+    def get_active_key_info(self, client_key_id):
+        ''' Get the info in the KeyRecord for the active secret in the DRM. '''
+        url = self.keyURL + '/active/' + urllib.quote_plus(client_key_id)
+        response = self.connection.get(url, headers=self.headers)
+        return KeyInfo.from_json(response.json())
+
+    @pki.handle_exceptions()
+    def modify_key_status(self, key_id, status):
+        ''' Modify the status of a key '''
+        url = self.keyURL + '/' + key_id
+        params = {'status':status}
+        self.connection.post(url, None, headers=self.headers, params=params)
+
+    @pki.handle_exceptions()
+    def approve_request(self, request_id):
+        ''' Approve a secret recovery request '''
+        url = self.keyRequestsURL + '/' + request_id + '/approve'
+        self.connection.post(url, self.headers)
+
+    @pki.handle_exceptions()
+    def reject_request(self, request_id):
+        ''' Reject a secret recovery request. '''
+        url = self.keyRequestsURL + '/' + request_id + '/reject'
+        self.connection.post(url, self.headers)
+
+    @pki.handle_exceptions()
+    def cancel_request(self, request_id):
+        ''' Cancel a secret recovery request '''
+        url = self.keyRequestsURL + '/' + request_id + '/cancel'
+        self.connection.post(url, self.headers)
+
+    def generate_pki_archive_options(self, trans_wrapped_session_key, session_wrapped_secret):
+        ''' Return a PKIArchiveOptions structure for archiving a secret
+
+            Takes in a session key wrapped by the DRM transport certificate,
+            and a secret wrapped with the session key and creates a PKIArchiveOptions
+            structure to be used when archiving a secret
+        '''
+        pass
+
+    def generate_archive_options(self, secret):
+        ''' Return a PKIArchiveOptions structure for archiving a secret.
+
+            This method uses NSS calls to do the following:
+            1) generate a session key
+            2) wrap the session key with the transport key
+            3) wrap the secret with the session key
+            4) create the PKIArchiveOptions structure using the results of
+               (2) and (3)
+
+            This method expects initialize_nss() to have been called previously.
+        '''
+        session_key = self.crypto.generate_symmetric_key()
+        trans_wrapped_session_key = self.crypto.asymmetric_wrap(session_key, self.transport_cert)
+        wrapped_secret = self.crypto.symmetric_wrap(secret, session_key)
+
+        return self.generate_pki_archive_options(trans_wrapped_session_key, wrapped_secret)
+
+    @pki.handle_exceptions()
     def create_request(self, request):
         ''' Submit an archival, recovery or key generation request
             to the DRM.
@@ -396,62 +425,55 @@ class KeyClient(object):
         return KeyRequestResponse.from_json(response.json())
 
     @pki.handle_exceptions()
-    def approve_request(self, request_id):
-        ''' Approve a secret recovery request '''
-        url = self.keyRequestsURL + '/' + request_id + '/approve'
-        self.connection.post(url, self.headers)
+    def generate_symmetric_key(self, client_key_id, algorithm, size, usages):
+        ''' Generate and archive a symmetric key on the DRM.
 
-    @pki.handle_exceptions()
-    def reject_request(self, request_id):
-        ''' Reject a secret recovery request. '''
-        url = self.keyRequestsURL + '/' + request_id + '/reject'
-        self.connection.post(url, self.headers)
-
-    @pki.handle_exceptions()
-    def cancel_request(self, request_id):
-        ''' Cancel a secret recovery request '''
-        url = self.keyRequestsURL + '/' + request_id + '/cancel'
-        self.connection.post(url, self.headers)
-
-    @pki.handle_exceptions()
-    def request_recovery(self, key_id, request_id=None, session_wrapped_passphrase=None,
-                        trans_wrapped_session_key=None, b64certificate=None, nonce_data=None):
-        ''' Create a request to recover a secret.
-
-            To retrieve a symmetric key or passphrase, the only parameter that is required is
-            the keyId.  It is possible (but not required) to pass in the session keys/passphrase
-            and nonceData for the retrieval at this time.  Those parameters are documented
-            in the docstring for retrieve_key below.
-
-            To retrieve an asymmetric key, the keyId and the the base-64 encoded certificate
-            is required.
+            Return a KeyRequestResponse which contains a KeyRequestInfo
+            object that describes the URL for the request and generated key.
         '''
-        request = KeyRecoveryRequest(key_id=key_id,
-                                     request_id=request_id,
-                                     trans_wrapped_session_key=trans_wrapped_session_key,
-                                     session_wrapped_passphrase=session_wrapped_passphrase,
-                                     certificate=b64certificate,
-                                     nonce_data=nonce_data)
+        request = SymKeyGenerationRequest(client_key_id=client_key_id,
+                                          key_size=size,
+                                          key_algorithm=algorithm,
+                                          key_usages=usages)
         return self.create_request(request)
 
     @pki.handle_exceptions()
-    def request_archival(self, client_key_id, data_type, wrapped_private_data,
+    def archive_key(self, client_key_id, data_type, private_data=None,
+                    wrapped_private_data=None,
                     key_algorithm=None, key_size=None):
         ''' Archive a secret (symmetric key or passphrase) on the DRM.
 
             Requires a user-supplied client ID.  There can be only one active
             key with a specified client ID.  If a record for a duplicate active
-            key exists, an exception is thrown.
+            key exists, a BadRequestException is thrown.
 
             data_type can be one of the following:
+                KeyRequestResource.SYMMETRIC_KEY_TYPE,
+                KeyRequestResource.ASYMMETRIC_KEY_TYPE,
+                KeyRequestResource.PASS_PHRASE_TYPE
+
+            key_algorithm and key_size are applicable to symmetric keys only.
+            If a symmetric key is being archived, these parameters are required.
 
             wrapped_private_data consists of a PKIArchiveOptions structure, which
             can be constructed using either generate_archive_options() or
             generate_pki_archive_options() below.
 
-            key_algorithm and key_size are applicable to symmetric keys only.
-            If a symmetric key is being archived, these parameters are required.
+            private_data is the secret that is to be archived.
+
+            Callers must specify EITHER wrapped_private_data OR private_data.
+            If wrapped_private_data is specified, then this data is forwarded to the
+            DRM unchanged.  Otherwise, the private_data is converted to a
+            PKIArchiveOptions structure using the functions below.
+
+            The function returns a KeyRequestResponse object containing a KeyRequestInfo
+            object with details about the archival request and key archived.
         '''
+        if wrapped_private_data == None:
+            if private_data == None:
+                raise ValueError("No data provided to be archived")
+            wrapped_private_data = self.generate_archive_options(private_data)
+
         request = KeyArchivalRequest(client_key_id=client_key_id,
                                      data_type=data_type,
                                      wrapped_private_data=wrapped_private_data,
@@ -460,26 +482,162 @@ class KeyClient(object):
         return self.create_request(request)
 
     @pki.handle_exceptions()
-    def get_key_info(self, key_id):
-        ''' Get the info in the KeyRecord for a specific secret in the DRM. '''
-        url = self.keyURL + '/' + key_id
-        response = self.connection.get(url, headers=self.headers)
-        return KeyInfo.from_json(response.json())
+    def recover_key(self, key_id, request_id=None, session_wrapped_passphrase=None,
+                        trans_wrapped_session_key=None, b64certificate=None, nonce_data=None):
+        ''' Create a request to recover a secret.
+
+            To retrieve a symmetric key or passphrase, the only parameter that is required is
+            the keyId.  It is possible (but not required) to pass in the session keys/passphrase
+            and nonceData for the retrieval at this time.  Those parameters are documented
+            in the docstring for retrieve_key below.
+
+            To retrieve an asymmetric key, the keyId and the the base-64 encoded certificate
+            is required.
+        '''
+        request = KeyRecoveryRequest(key_id=key_id,
+                                     request_id=request_id,
+                                     trans_wrapped_session_key=trans_wrapped_session_key,
+                                     session_wrapped_passphrase=session_wrapped_passphrase,
+                                     certificate=b64certificate,
+                                     nonce_data=nonce_data)
+        return self.create_request(request)
 
     @pki.handle_exceptions()
-    def get_active_key_info(self, client_key_id):
-        ''' Get the info in the KeyRecord for the active secret in the DRM. '''
-        url = self.keyURL + '/active/' + urllib.quote_plus(client_key_id)
-        response = self.connection.get(url, headers=self.headers)
-        print response
-        return KeyInfo.from_json(response.json())
+    def retrieve_key_data(self, data):
+        ''' Retrieve a secret from the DRM.
+
+            @param: data - a KeyRecoveryRequest containing the keyId of the
+            secret being retrieved, the request_id of the approved recovery
+            request and a wrapping mechanism.  More details at
+            KRAClient.retrieve_key.
+
+            Returns a KeyData object containing the wrapped secret.
+        '''
+        url = self.keyURL + '/retrieve'
+        keyRequest = json.dumps(data, cls=encoder.CustomTypeEncoder, sort_keys=True)
+        response = self.connection.post(url, keyRequest, self.headers)
+        return KeyData.from_json(response.json())
 
     @pki.handle_exceptions()
-    def modify_key_status(self, key_id, status):
-        ''' Modify the status of a key '''
-        url = self.keyURL + '/' + key_id
-        params = {'status':status}
-        self.connection.post(url, None, headers=self.headers, params=params)
+    def retrieve_key(self, key_id, trans_wrapped_session_key=None):
+        ''' Retrieve a secret (passphrase or symmetric key) from the DRM.
+
+        This function generates a key recovery request, approves it, and retrieves
+        the secret referred to by key_id.  This assumes that only one approval is required
+        to authorize the recovery.
+
+        To ensure data security in transit, the data will be returned encrypted by a session
+        key (56 bit DES3 symmetric key) - which is first wrapped (encrypted) by the public
+        key of the DRM transport certificate before being sent to the DRM.  The
+        parameter trans_wrapped_session_key refers to this wrapped session key.
+
+        There are two ways of using this function:
+
+        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)
+
+        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
+        wrapped secret and some nonce data to be used as a salt when unwrapping.
+        '''
+        key_provided = True
+        if (trans_wrapped_session_key == None):
+            key_provided = False
+            session_key = self.crypto.generate_symmetric_key()
+            trans_wrapped_session_key = self.crypto.asymmetric_wrap(session_key,
+                                                                    self.transport_cert)
+
+        response = self.recover_key(key_id)
+        request_id = response.get_request_id()
+        self.approve_request(request_id)
+
+        request = KeyRecoveryRequest(
+                        key_id=key_id,
+                        request_id=request_id,
+                        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(
+                                base64.decodestring(key_data.wrappedPrivateData),
+                                session_key,
+                                nonce_iv=base64.decodestring(key_data.nonceData))
+        return key_data, unwrapped_key
+
+    @pki.handle_exceptions()
+    def retrieve_key_by_passphrase(self, key_id, passphrase=None,
+                                   trans_wrapped_session_key=None,
+                                   session_wrapped_passphrase=None,
+                                   nonce_data=None):
+        ''' Retrieve a secret (passphrase or symmetric key) from the DRM using a passphrase.
+
+        This function generates a key recovery request, approves it, and retrieves
+        the secret referred to by key_id.  This assumes that only one approval is required
+        to authorize the recovery.
+
+        The secret is secured in transit by wrapping the secret with a passphrase using
+        PBE encryption.
+
+        There are two ways of using this function:
+
+        1) A passphrase is provided by the caller.
+
+        In this case, CryptoUtil methods will be called to create the data to securely send the
+        passphrase to the DRM.  Basically, three pieces of data will be sent:
+
+        - the passphrase wrapped by a 56 bit DES3 symmetric key (the session key).  This
+          is referred to as the parameter session_wrapped_passphrase above.
+
+        - the session key wrapped with the public key in the DRM transport certificate.  This
+          is referred to as the trans_wrapped_session_key above.
+
+        - ivps nonce data, referred to as nonce_data
+
+        The function will return the tuple (KeyData, unwrapped_secret)
+
+        2) The caller provides the trans_wrapped_session_key, session_wrapped_passphrase
+        and nonce_data.
+
+        In this case, the data will simply be passed to the DRM.  The function will return
+        the secret encrypted by the passphrase using PBE Encryption.  The secret will still
+        need to be decrypted by the caller.
+
+        The function will return the tuple (KeyData, None)
+        '''
+        pass
+
+    @pki.handle_exceptions()
+    def retrieve_key_by_pkcs12(self, key_id, certificate, passphrase):
+        ''' Retrieve an asymmetric private key and return it as PKCS12 data.
+
+        This function generates a key recovery request, approves it, and retrieves
+        the secret referred to by key_id in a PKCS12 file.  This assumes that only
+        one approval is required to authorize the recovery.
+
+        This function requires the following parameters:
+        - key_id : the ID of the key
+        - certificate: the certificate associated with the private key
+        - passphrase: A passphrase for the pkcs12 file.
+
+        The function returns a KeyData object.
+        '''
+        response = self.recover_key(key_id, b64certificate=certificate)
+        request_id = response.get_request_id()
+        self.approve_request(request_id)
+
+        request = KeyRecoveryRequest(key_id=key_id,
+                                     request_id=request_id,
+                                     passphrase=passphrase)
+
+        return self.retrieve_key_data(request)
 
 encoder.NOTYPES['Attribute'] = pki.Attribute
 encoder.NOTYPES['AttributeList'] = pki.AttributeList
diff --git a/base/common/python/pki/kraclient.py b/base/common/python/pki/kraclient.py
index 25c4dc9ca57e659c2440baba3b1d5d616ee1b5ee..4d62f38f9af63043debdedcde703b4ad83ae2418 100644
--- a/base/common/python/pki/kraclient.py
+++ b/base/common/python/pki/kraclient.py
@@ -25,7 +25,6 @@ to interact with the DRM to expose the functionality of the KeyClient and
 KeyRequestResouce REST APIs.
 '''
 
-import base64
 import pki.key as key
 
 from pki.systemcert import SystemCertClient
@@ -49,204 +48,8 @@ class KRAClient(object):
                         beforehand.
         '''
         self.connection = connection
-        self.keys = key.KeyClient(connection)
+        self.keys = key.KeyClient(connection, crypto, transport_cert_nick)
         self.system_certs = SystemCertClient(connection)
-        self.crypto = crypto
-        if transport_cert_nick != None:
-            self.transport_cert = crypto.get_cert(transport_cert_nick)
-        else:
-            self.transport_cert = None
 
-    def set_transport_cert(self, transport_cert_nick):
-        ''' Set the transport certificate for crypto operations '''
-        self.transport_cert = self.crypto.get_cert(transport_cert_nick)
 
-    def retrieve_key(self, key_id, trans_wrapped_session_key=None):
-        ''' Retrieve a secret (passphrase or symmetric key) from the DRM.
-
-        This function generates a key recovery request, approves it, and retrieves
-        the secret referred to by key_id.  This assumes that only one approval is required
-        to authorize the recovery.
-
-        To ensure data security in transit, the data will be returned encrypted by a session
-        key (56 bit DES3 symmetric key) - which is first wrapped (encrypted) by the public
-        key of the DRM transport certificate before being sent to the DRM.  The
-        parameter trans_wrapped_session_key refers to this wrapped session key.
-
-        There are two ways of using this function:
-
-        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)
-
-        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
-        wrapped secret and some nonce data to be used as a salt when unwrapping.
-        '''
-        key_provided = True
-        if (trans_wrapped_session_key == None):
-            key_provided = False
-            session_key = self.crypto.generate_symmetric_key()
-            trans_wrapped_session_key = self.crypto.asymmetric_wrap(session_key,
-                                                                    self.transport_cert)
-
-        response = self.keys.request_recovery(key_id)
-        request_id = response.get_request_id()
-        self.keys.approve_request(request_id)
-
-        key_data = self.keys.request_key_retrieval(key_id, request_id,
-                        trans_wrapped_session_key=base64.encodestring(trans_wrapped_session_key))
-        if key_provided:
-            return key_data, None
-
-        unwrapped_key = self.crypto.symmetric_unwrap(
-                                base64.decodestring(key_data.wrappedPrivateData),
-                                session_key,
-                                nonce_iv=base64.decodestring(key_data.nonceData))
-        return key_data, unwrapped_key
-
-    def retrieve_key_by_passphrase(self, key_id, passphrase=None,
-                                   trans_wrapped_session_key=None,
-                                   session_wrapped_passphrase=None,
-                                   nonce_data=None):
-        ''' Retrieve a secret (passphrase or symmetric key) from the DRM using a passphrase.
-
-        This function generates a key recovery request, approves it, and retrieves
-        the secret referred to by key_id.  This assumes that only one approval is required
-        to authorize the recovery.
-
-        The secret is secured in transit by wrapping the secret with a passphrase using
-        PBE encryption.
-
-        There are two ways of using this function:
-
-        1) A passphrase is provided by the caller.
-
-        In this case, CryptoUtil methods will be called to create the data to securely send the
-        passphrase to the DRM.  Basically, three pieces of data will be sent:
-
-        - the passphrase wrapped by a 56 bit DES3 symmetric key (the session key).  This
-          is referred to as the parameter session_wrapped_passphrase above.
-
-        - the session key wrapped with the public key in the DRM transport certificate.  This
-          is referred to as the trans_wrapped_session_key above.
-
-        - ivps nonce data, referred to as nonce_data
-
-        The function will return the tuple (KeyData, unwrapped_secret)
-
-        2) The caller provides the trans_wrapped_session_key, session_wrapped_passphrase
-        and nonce_data.
-
-        In this case, the data will simply be passed to the DRM.  The function will return
-        the secret encrypted by the passphrase using PBE Encryption.  The secret will still
-        need to be decrypted by the caller.
-
-        The function will return the tuple (KeyData, None)
-        '''
-        pass
-
-    def retrieve_key_by_pkcs12(self, key_id, certificate, passphrase):
-        ''' Retrieve an asymmetric private key and return it as PKCS12 data.
-
-        This function generates a key recovery request, approves it, and retrieves
-        the secret referred to by key_id in a PKCS12 file.  This assumes that only
-        one approval is required to authorize the recovery.
-
-        This function requires the following parameters:
-        - key_id : the ID of the key
-        - certificate: the certificate associated with the private key
-        - passphrase: A passphrase for the pkcs12 file.
-
-        The function returns a KeyData object.
-        '''
-        response = self.keys.request_recovery(key_id, b64certificate=certificate)
-        request_id = response.get_request_id()
-        self.keys.approve_request(request_id)
-
-        return self.keys.request_key_retrieval(key_id, request_id, passphrase)
-
-
-    def generate_symmetric_key(self, client_key_id, algorithm, size, usages):
-        ''' Generate and archive a symmetric key on the DRM.
-
-            Return a KeyRequestResponse which contains a KeyRequestInfo
-            object that describes the URL for the request and generated key.
-        '''
-        request = key.SymKeyGenerationRequest(client_key_id=client_key_id,
-                                              key_size=size,
-                                              key_algorithm=algorithm,
-                                              key_usages=usages)
-        return self.keys.create_request(request)
-
-    def archive_key(self, client_key_id, data_type, private_data=None,
-                    wrapped_private_data=None,
-                    key_algorithm=None, key_size=None):
-        ''' Archive a secret (symmetric key or passphrase) on the DRM.
-
-            Requires a user-supplied client ID.  There can be only one active
-            key with a specified client ID.  If a record for a duplicate active
-            key exists, a BadRequestException is thrown.
-
-            data_type can be one of the following:
-                KeyRequestResource.SYMMETRIC_KEY_TYPE,
-                KeyRequestResource.ASYMMETRIC_KEY_TYPE,
-                KeyRequestResource.PASS_PHRASE_TYPE
-
-            key_algorithm and key_size are applicable to symmetric keys only.
-            If a symmetric key is being archived, these parameters are required.
-
-            wrapped_private_data consists of a PKIArchiveOptions structure, which
-            can be constructed using either generate_archive_options() or
-            generate_pki_archive_options() below.
-
-            private_data is the secret that is to be archived.
-
-            Callers must specify EITHER wrapped_private_data OR private_data.
-            If wrapped_private_data is specified, then this data is forwarded to the
-            DRM unchanged.  Otherwise, the private_data is converted to a
-            PKIArchiveOptions structure using the functions below.
-
-            The function returns a KeyRequestResponse object containing a KeyRequestInfo
-            object with details about the archival request and key archived.
-        '''
-        if wrapped_private_data == None:
-            if private_data == None:
-                # raise BadRequestException - to be added in next patch
-                return None
-            wrapped_private_data = self.generate_archive_options(private_data)
-        return self.keys.request_archival(client_key_id, data_type, wrapped_private_data,
-                                          key_algorithm, key_size)
-
-    def generate_pki_archive_options(self, trans_wrapped_session_key, session_wrapped_secret):
-        ''' Return a PKIArchiveOptions structure for archiving a secret
-
-            Takes in a session key wrapped by the DRM transport certificate,
-            and a secret wrapped with the session key and creates a PKIArchiveOptions
-            structure to be used when archiving a secret
-        '''
-        pass
-
-    def generate_archive_options(self, secret):
-        ''' Return a PKIArchiveOptions structure for archiving a secret.
-
-            This method uses NSS calls to do the following:
-            1) generate a session key
-            2) wrap the session key with the transport key
-            3) wrap the secret with the session key
-            4) create the PKIArchiveOptions structure using the results of
-               (2) and (3)
-
-            This method expects initialize_nss() to have been called previously.
-        '''
-        session_key = self.crypto.generate_symmetric_key()
-        trans_wrapped_session_key = self.crypto.asymmetric_wrap(session_key, self.transport_cert)
-        wrapped_secret = self.crypto.symmetric_wrap(secret, session_key)
-
-        return self.generate_pki_archive_options(trans_wrapped_session_key, wrapped_secret)
 
diff --git a/base/kra/functional/drmtest.py b/base/kra/functional/drmtest.py
index 90cc8b2c8fdb122b51632bbab53caf49bf247247..a1db60c55ba0cfd6c4029b19affa5558ea3932bc 100644
--- a/base/kra/functional/drmtest.py
+++ b/base/kra/functional/drmtest.py
@@ -63,31 +63,42 @@ def print_key_data(key_data):
 
 def main():
     ''' test code execution '''
+
+    # set up the connectoon to the DRM, including authentication credentials
     connection = PKIConnection('https', 'localhost', '8443', 'kra')
     connection.set_authentication_cert('/tmp/temp4.pem')
+
+    # create an NSS DB for crypto operations
     certdb_dir = "/tmp/drmtest-certdb"
     certdb_password = "redhat123"
-    transport_nick = "kra transport cert"
     cryptoutil.NSSCryptoUtil.setup_database(certdb_dir, certdb_password, over_write=True)
+
+    #create kraclient
     crypto = cryptoutil.NSSCryptoUtil(certdb_dir, certdb_password)
     kraclient = KRAClient(connection, crypto)
+    keyclient = kraclient.keys
 
-    # Test 1: Get transport certificate and import it into the NSS database
+    # Get transport cert and insert in the certdb
+    transport_nick = "kra transport cert"
     transport_cert = kraclient.system_certs.get_transport_cert()
-    print transport_cert
     tcert = transport_cert[len(pki.CERT_HEADER):len(transport_cert) -len(pki.CERT_FOOTER)]
     crypto.import_cert(transport_nick, base64.decodestring(tcert), "u,u,u")
-    crypto.initialize_db()
-    kraclient.set_transport_cert(transport_nick)
+
+    # initialize the certdb for crypto operations
+    # for NSS db, this must be done after importing the transport cert
+    crypto.initialize()
+
+    # set transport cert into keyclient
+    keyclient.set_transport_cert(transport_nick)
 
     # Test 2: Get key request info
     print "Now getting key request"
-    keyrequest = kraclient.keys.get_request_info('2')
+    keyrequest = keyclient.get_request_info('2')
     print_key_request(keyrequest)
 
     # Test 3: List requests
     print "Now listing some requests"
-    keyrequests = kraclient.keys.list_requests('complete', 'securityDataRecovery')
+    keyrequests = keyclient.list_requests('complete', 'securityDataRecovery')
     print keyrequests.key_requests
     for request in keyrequests.key_requests:
         print_key_request(request)
@@ -95,18 +106,18 @@ def main():
     # Test 4: generate symkey -- same as barbican_encode()
     print "Now generating symkey on KRA"
     #client_key_id = "Vek #1" + time.strftime('%X %x %Z')
-    client_key_id = "abcxyz"
+    client_key_id = "vek1234567"
     algorithm = "AES"
     key_size = 128
     usages = [key.SymKeyGenerationRequest.DECRYPT_USAGE, key.SymKeyGenerationRequest.ENCRYPT_USAGE]
-    response = kraclient.generate_symmetric_key(client_key_id, algorithm, key_size, usages)
+    response = keyclient.generate_symmetric_key(client_key_id, algorithm, key_size, usages)
     print_key_request(response.requestInfo)
     print "Request ID is " + response.requestInfo.get_request_id()
     key_id = response.get_key_id()
 
     # Test 5: Confirm the key_id matches
     print "Now getting key ID for clientKeyID=\"" + client_key_id + "\""
-    key_infos = kraclient.keys.list_keys(client_key_id=client_key_id, status="active")
+    key_infos = keyclient.list_keys(client_key_id=client_key_id, status="active")
     for key_info in key_infos.key_infos:
         print_key_info(key_info)
         key_id2 = key_info.get_key_id()
@@ -117,8 +128,8 @@ def main():
 
     # Test 6: Barbican_decode() - Retrieve while providing trans_wrapped_session_key
     session_key = crypto.generate_symmetric_key()
-    wrapped_session_key = crypto.asymmetric_wrap(session_key, kraclient.transport_cert)
-    key_data, _unwrapped_key = kraclient.retrieve_key(key_id, trans_wrapped_session_key=wrapped_session_key)
+    wrapped_session_key = crypto.asymmetric_wrap(session_key, keyclient.transport_cert)
+    key_data, _unwrapped_key = keyclient.retrieve_key(key_id, trans_wrapped_session_key=wrapped_session_key)
     print_key_data(key_data)
     unwrapped_key = crypto.symmetric_unwrap(base64.decodestring(key_data.wrappedPrivateData),
                                             session_key,
@@ -126,7 +137,7 @@ def main():
     key1 = base64.encodestring(unwrapped_key)
 
     # Test 7: Recover key without providing trans_wrapped_session_key
-    key_data, unwrapped_key = kraclient.retrieve_key(key_id)
+    key_data, unwrapped_key = keyclient.retrieve_key(key_id)
     print_key_data(key_data)
     key2 = base64.encodestring(unwrapped_key)
 
@@ -141,21 +152,21 @@ def main():
     # Test 10 = test BadRequestException on create()
     print "Trying to generate a new symkey with the same client ID"
     try:
-        response = kraclient.generate_symmetric_key(client_key_id, algorithm, key_size, usages)
+        response = keyclient.generate_symmetric_key(client_key_id, algorithm, key_size, usages)
     except pki.BadRequestException as exc:
         print "BadRequestException thrown - Code:" + exc.code + " Message: " + exc.message
 
     # Test 11 - Test RequestNotFoundException on get_request_info
     print "Try to list a nonexistent request"
     try:
-        keyrequest = kraclient.keys.get_request_info('200000034')
+        keyrequest = keyclient.get_request_info('200000034')
     except pki.RequestNotFoundException as exc:
         print "RequestNotFoundException thrown - Code:" + exc.code + " Message: " + exc.message
 
     # Test 12 - Test exception on retrieve_key.
     print "Try to retrieve an invalid key"
     try:
-        key_data, unwrapped_key = kraclient.retrieve_key('2000003434')
+        key_data, unwrapped_key = keyclient.retrieve_key('2000003434')
     except pki.KeyNotFoundException as exc:
         print "KeyNotFoundException thrown - Code:" + exc.code + " Message: " + exc.message
     except pki.PKIException as exc:
@@ -165,30 +176,30 @@ def main():
 
     #Test 13 = getKeyInfo
     print "Get key info for existing key"
-    key_info = kraclient.keys.get_key_info(key_id)
+    key_info = keyclient.get_key_info(key_id)
     print_key_info(key_info)
 
     # Test 14: get the active key
     print "Get the active key for client id: " + client_key_id
-    key_info = kraclient.keys.get_active_key_info(client_key_id)
+    key_info = keyclient.get_active_key_info(client_key_id)
     print_key_info(key_info)
 
     #Test 15: change the key status
     print "Change the key status"
-    kraclient.keys.modify_key_status(key_id, "inactive")
-    print_key_info(kraclient.keys.get_key_info(key_id))
+    keyclient.modify_key_status(key_id, "inactive")
+    print_key_info(keyclient.get_key_info(key_id))
 
     # Test 16: Get key info for non-existent key
     print "Get key info for non-existent key"
     try:
-        key_info = kraclient.keys.get_key_info('200004556')
+        key_info = keyclient.get_key_info('200004556')
     except pki.KeyNotFoundException as exc:
         print "KeyNotFoundException thrown - Code:" + exc.code + " Message: " + exc.message
 
     # Test 17: Get key info for non-existent active key
     print "Get non-existent active key"
     try:
-        key_info = kraclient.keys.get_active_key_info(client_key_id)
+        key_info = keyclient.get_active_key_info(client_key_id)
     except pki.ResourceNotFoundException as exc:
         print "ResourceNotFoundException thrown - Code: " + exc.code + "Message: " + exc.message
 
-- 
1.8.4.2

