From 4c59f4fb724d5297f6859126221ef6d4a5273dba Mon Sep 17 00:00:00 2001
From: Hayden Roche <haydenroche5@gmail.com>
Date: Mon, 24 Jan 2022 13:38:30 -0800
Subject: [PATCH] Call wolfCrypt_SetPrivateKeyReadEnable_fips at init, if
 necessary.

This is another FIPS v5 change. This function needs to be called before we can
do private key ops with the new FIPS module.
---
 src/wolfcrypt/__init__.py   |  7 +++++++
 src/wolfcrypt/_build_ffi.py | 10 ++++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/wolfcrypt/__init__.py b/src/wolfcrypt/__init__.py
index d4c57ab..868a154 100644
--- a/src/wolfcrypt/__init__.py
+++ b/src/wolfcrypt/__init__.py
@@ -60,3 +60,10 @@ else:
             ret = _lib.wc_SetSeed_Cb(_ffi.addressof(_lib, "wc_GenerateSeed"))
             if ret < 0:
                 raise WolfCryptError("wc_SetSeed_Cb failed (%d)" % ret)
+    if _lib.FIPS_ENABLED and (_lib.FIPS_VERSION > 5 or (_lib.FIPS_VERSION == 5
+        and _lib.FIPS_VERSION >= 1)):
+        ret = _lib.wolfCrypt_SetPrivateKeyReadEnable_fips(1,
+                                                          _lib.WC_KEYTYPE_ALL);
+        if ret < 0:
+            raise WolfCryptError("wolfCrypt_SetPrivateKeyReadEnable_fips failed"
+                " (%d)" % ret)
diff --git a/src/wolfcrypt/_build_ffi.py b/src/wolfcrypt/_build_ffi.py
index d706b63..d28b9ab 100644
--- a/src/wolfcrypt/_build_ffi.py
+++ b/src/wolfcrypt/_build_ffi.py
@@ -499,6 +499,16 @@ if WC_RNG_SEED_CB_ENABLED:
     int wc_SetSeed_Cb(wc_RngSeed_Cb cb);
     """
 
+if FIPS_ENABLED and (FIPS_VERSION > 5 or (FIPS_VERSION == 5 and FIPS_VERSION >= 1)):
+    _cdef += """
+    enum wc_KeyType {
+        WC_KEYTYPE_ALL = 0
+    };
+
+    int wolfCrypt_SetPrivateKeyReadEnable_fips(int, enum wc_KeyType);
+    int wolfCrypt_GetPrivateKeyReadEnable_fips(enum wc_KeyType);
+    """
+
 ffibuilder.cdef(_cdef)
 
 if __name__ == "__main__":