From ea8893cd179814f296aa73baa9c1aa04dcd86e98 Mon Sep 17 00:00:00 2001
From: deatil <2217957370@qq.com>
Date: Thu, 8 Aug 2024 10:32:12 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../go-cryptobin/pkcs8/pbes2/pkcs8.go         |  8 ++--
 pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8.go    | 35 ++++++++------
 .../go-cryptobin/pkcs8/pkcs8_test.go          | 46 ++++++++++++++++++-
 3 files changed, 71 insertions(+), 18 deletions(-)

diff --git a/pkg/lakego-pkg/go-cryptobin/pkcs8/pbes2/pkcs8.go b/pkg/lakego-pkg/go-cryptobin/pkcs8/pbes2/pkcs8.go
index 52400347..f9995290 100644
--- a/pkg/lakego-pkg/go-cryptobin/pkcs8/pbes2/pkcs8.go
+++ b/pkg/lakego-pkg/go-cryptobin/pkcs8/pbes2/pkcs8.go
@@ -73,11 +73,11 @@ type pbes2Params struct {
 
 // 加密 PKCS8
 func EncryptPKCS8PrivateKey(
-    rand io.Reader,
+    rand      io.Reader,
     blockType string,
-    data []byte,
-    password []byte,
-    opts ...Opts,
+    data      []byte,
+    password  []byte,
+    opts      ...Opts,
 ) (*pem.Block, error) {
     opt := &DefaultOpts
     if len(opts) > 0 {
diff --git a/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8.go b/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8.go
index 9e6ca902..d4475c46 100644
--- a/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8.go
+++ b/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8.go
@@ -135,9 +135,10 @@ var (
 
 type (
     // 配置
-    Opts       = pbes2.Opts
-    PBKDF2Opts = pbes2.PBKDF2Opts
-    ScryptOpts = pbes2.ScryptOpts
+    Opts         = pbes2.Opts
+    PBKDF2Opts   = pbes2.PBKDF2Opts
+    SMPBKDF2Opts = pbes2.SMPBKDF2Opts
+    ScryptOpts   = pbes2.ScryptOpts
 )
 
 var (
@@ -147,20 +148,28 @@ var (
     GetHashFromName   = pbes2.GetHashFromName
 )
 
-// 默认配置 PBKDF2
-var DefaultPBKDF2Opts = pbes2.DefaultPBKDF2Opts
+var (
+    // 默认 Hash
+    DefaultHash   = pbes2.DefaultHash
+    DefaultSMHash = pbes2.DefaultSMHash
+)
 
-// 默认配置 GmSM PBKDF2
-var DefaultSMPBKDF2Opts = pbes2.DefaultSMPBKDF2Opts
+var (
+    // 默认配置 PBKDF2
+    DefaultPBKDF2Opts = pbes2.DefaultPBKDF2Opts
 
-// 默认配置 Scrypt
-var DefaultScryptOpts = pbes2.DefaultScryptOpts
+    // 默认配置 GmSM PBKDF2
+    DefaultSMPBKDF2Opts = pbes2.DefaultSMPBKDF2Opts
 
-// 默认配置
-var DefaultOpts = pbes2.DefaultOpts
+    // 默认配置 Scrypt
+    DefaultScryptOpts = pbes2.DefaultScryptOpts
 
-// 默认 GmSM 配置
-var DefaultSMOpts = pbes2.DefaultSMOpts
+    // 默认配置
+    DefaultOpts = pbes2.DefaultOpts
+
+    // 默认 GmSM 配置
+    DefaultSMOpts = pbes2.DefaultSMOpts
+)
 
 // 解析设置
 // opt, err := ParseOpts("AES256CBC", "SHA256")
diff --git a/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8_test.go b/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8_test.go
index 68b5c60d..f2548215 100644
--- a/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8_test.go
+++ b/pkg/lakego-pkg/go-cryptobin/pkcs8/pkcs8_test.go
@@ -1,6 +1,7 @@
 package pkcs8
 
 import (
+    "bytes"
     "testing"
     "crypto/rsa"
     "crypto/rand"
@@ -1105,6 +1106,40 @@ func Test_EncryptPEMBlock_Gost(t *testing.T) {
 }
 
 func Test_EncryptPEMBlock_GmSMOpts(t *testing.T) {
+    t.Run("DefaultSMOpts", func(t *testing.T) {
+        test_EncryptPEMBlock_GmSMOpts(t, DefaultSMOpts)
+    })
+
+    t.Run("DefaultSMPBKDF2Opts SM4CFB", func(t *testing.T) {
+        test_EncryptPEMBlock_GmSMOpts(t, Opts{
+            Cipher:  SM4CFB,
+            KDFOpts: DefaultSMPBKDF2Opts,
+        })
+    })
+
+    t.Run("SMPBKDF2Opts SM4CFB", func(t *testing.T) {
+        test_EncryptPEMBlock_GmSMOpts(t, Opts{
+            Cipher:  SM4CFB,
+            KDFOpts: SMPBKDF2Opts{
+                SaltSize:       8,
+                IterationCount: 5000,
+                HMACHash:       DefaultSMHash,
+            },
+        })
+    })
+
+    t.Run("SMPBKDF2Opts no HMACHash", func(t *testing.T) {
+        test_EncryptPEMBlock_GmSMOpts(t, Opts{
+            Cipher:  SM4CFB,
+            KDFOpts: SMPBKDF2Opts{
+                SaltSize:       8,
+                IterationCount: 6000,
+            },
+        })
+    })
+}
+
+func test_EncryptPEMBlock_GmSMOpts(t *testing.T, opts Opts) {
     block, _ := pem.Decode([]byte(testKey_des_EDE3_CBC))
 
     bys, err := DecryptPEMBlock(block, []byte("123"))
@@ -1112,7 +1147,7 @@ func Test_EncryptPEMBlock_GmSMOpts(t *testing.T) {
         t.Fatal("PEM data decrypted error: " + err.Error())
     }
 
-    enblock, err := EncryptPEMBlock(rand.Reader, "ENCRYPTED PRIVATE KEY", bys, []byte("test-passsss"), DefaultSMOpts)
+    enblock, err := EncryptPEMBlock(rand.Reader, "ENCRYPTED PRIVATE KEY", bys, []byte("test-passsss"), opts)
     if err != nil {
         t.Error("encrypt: ", err)
     }
@@ -1124,4 +1159,13 @@ func Test_EncryptPEMBlock_GmSMOpts(t *testing.T) {
     if enblock.Type != "ENCRYPTED PRIVATE KEY" {
         t.Errorf("unexpected enblock type; got %q want %q", enblock.Type, "RSA PRIVATE KEY")
     }
+
+    bys2, err := DecryptPEMBlock(enblock, []byte("test-passsss"))
+    if err != nil {
+        t.Fatal("data decrypted error: " + err.Error())
+    }
+
+    if bytes.Compare(bys2, bys) != 0 {
+        t.Errorf("DecryptPEMBlock error, got %x, want %x", bys2, bys)
+    }
 }