diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 6e80cec3d..55b7d2f5b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1800,6 +1800,7 @@ psa_status_t psa_import_key( const psa_key_attributes_t *attributes,
     psa_status_t status;
     psa_key_slot_t *slot = NULL;
     psa_se_drv_table_entry_t *driver = NULL;
+    size_t bits;
 
     *key = MBEDTLS_SVC_KEY_ID_INIT;
 
@@ -1814,72 +1815,33 @@ psa_status_t psa_import_key( const psa_key_attributes_t *attributes,
     if( status != PSA_SUCCESS )
         goto exit;
 
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( driver != NULL )
+    /* In the case of a transparent key or an opaque key stored in local
+     * storage (thus not in the case of generating a key in a secure element
+     * or cryptoprocessor with storage), we have to allocate a buffer to
+     * hold the generated key material. */
+    if( slot->key.data == NULL )
     {
-        const psa_drv_se_t *drv = psa_get_se_driver_methods( driver );
-        /* The driver should set the number of key bits, however in
-         * case it doesn't, we initialize bits to an invalid value. */
-        size_t bits = PSA_MAX_KEY_BITS + 1;
-        if( drv->key_management == NULL ||
-            drv->key_management->p_import == NULL )
-        {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto exit;
-        }
-        status = drv->key_management->p_import(
-            psa_get_se_driver_context( driver ),
-            psa_key_slot_get_slot_number( slot ),
-            attributes, data, data_length, &bits );
+        status = psa_allocate_buffer_to_slot( slot, data_length );
         if( status != PSA_SUCCESS )
             goto exit;
-        if( bits > PSA_MAX_KEY_BITS )
-        {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto exit;
-        }
-        slot->attr.bits = (psa_key_bits_t) bits;
     }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    if( psa_key_lifetime_is_external( psa_get_key_lifetime( attributes ) ) )
+
+    bits = slot->attr.bits;
+    status = psa_driver_wrapper_import_key( attributes,
+                                            data, data_length,
+                                            slot->key.data,
+                                            slot->key.bytes,
+                                            &slot->key.bytes, &bits );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( slot->attr.bits == 0 )
+        slot->attr.bits = (psa_key_bits_t) bits;
+    else if( bits != slot->attr.bits )
     {
-        /* Importing a key with external lifetime through the driver wrapper
-         * interface is not yet supported. Return as if this was an invalid
-         * lifetime. */
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }
-    else
-    {
-       /* In the case of a transparent key or an opaque key stored in local
-        * storage (thus not in the case of generating a key in a secure element
-        * or cryptoprocessor with storage), we have to allocate a buffer to
-        * hold the generated key material. */
-        if( slot->key.data == NULL )
-        {
-            status = psa_allocate_buffer_to_slot( slot, data_length );
-            if( status != PSA_SUCCESS )
-                goto exit;
-        }
-
-        size_t bits = slot->attr.bits;
-        status = psa_driver_wrapper_import_key( attributes,
-                                                data, data_length,
-                                                slot->key.data,
-                                                slot->key.bytes,
-                                                &slot->key.bytes, &bits );
-        if( status != PSA_SUCCESS )
-            goto exit;
-
-        if( slot->attr.bits == 0 )
-            slot->attr.bits = (psa_key_bits_t) bits;
-        else if( bits != slot->attr.bits )
-        {
-            status = PSA_ERROR_INVALID_ARGUMENT;
-            goto exit;
-        }
-    }
 
     status = psa_validate_optional_attributes( slot, attributes );
     if( status != PSA_SUCCESS )
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index fe54b7b97..fd03ed2bf 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -422,6 +422,35 @@ psa_status_t psa_driver_wrapper_import_key(
     psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
                                       psa_get_key_lifetime( attributes ) );
 
+    /* Try dynamically-registered SE interface first */
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    const psa_drv_se_t *drv;
+    psa_drv_se_context_t *drv_context;
+
+    if( psa_get_se_driver( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( drv->key_management == NULL ||
+            drv->key_management->p_import == NULL )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        /* The driver should set the number of key bits, however in
+         * case it doesn't, we initialize bits to an invalid value. */
+        *bits = PSA_MAX_KEY_BITS + 1;
+        status = drv->key_management->p_import(
+            drv_context,
+            *( (psa_key_slot_number_t *)key_buffer ),
+            attributes, data, data_length, bits );
+
+        if( status != PSA_SUCCESS )
+            return( status );
+
+        if( (*bits) > PSA_MAX_KEY_BITS )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        return( PSA_SUCCESS );
+    }
+#endif /* PSA_CRYPTO_SE_C */
+
     switch( location )
     {
         case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -445,9 +474,10 @@ psa_status_t psa_driver_wrapper_import_key(
                                               key_buffer_length, bits ) );
 
         default:
-            /* Key is declared with a lifetime not known to us */
+            /* Importing a key with external storage in not yet supported.
+             * Return in error indicating that the lifetime is not valid. */
             (void)status;
-            return( PSA_ERROR_NOT_SUPPORTED );
+            return( PSA_ERROR_INVALID_ARGUMENT );
     }
 
 }