From f3f50b4f6c3be5386fcffe55dc9448bc93e40934 Mon Sep 17 00:00:00 2001
From: starainrt <i@b612.me>
Date: Sun, 17 Mar 2024 17:29:46 +0800
Subject: [PATCH] update

---
 asy.go | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/asy.go b/asy.go
index ea72059..fe685c1 100644
--- a/asy.go
+++ b/asy.go
@@ -3,6 +3,7 @@ package starcrypto
 import (
 	"crypto"
 	"crypto/ecdsa"
+	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
@@ -17,7 +18,22 @@ func EncodePrivateKey(private crypto.PrivateKey, secret string) ([]byte, error)
 	case *ecdsa.PrivateKey:
 		return EncodeEcdsaPrivateKey(private.(*ecdsa.PrivateKey), secret)
 	default:
-		return nil, errors.New("private key type error")
+		b, err := x509.MarshalPKCS8PrivateKey(private)
+		if err != nil {
+			return nil, err
+		}
+		if secret == "" {
+			return pem.EncodeToMemory(&pem.Block{
+				Bytes: b,
+				Type:  "PRIVATE KEY",
+			}), err
+		}
+		chiper := x509.PEMCipherAES256
+		blk, err := x509.EncryptPEMBlock(rand.Reader, "PRIVATE KEY", b, []byte(secret), chiper)
+		if err != nil {
+			return nil, err
+		}
+		return pem.EncodeToMemory(blk), err
 	}
 }
 
@@ -28,7 +44,14 @@ func EncodePublicKey(public crypto.PublicKey) ([]byte, error) {
 	case *ecdsa.PublicKey:
 		return EncodeEcdsaPublicKey(public.(*ecdsa.PublicKey))
 	default:
-		return nil, errors.New("public key type error")
+		publicBytes, err := x509.MarshalPKIXPublicKey(public)
+		if err != nil {
+			return nil, err
+		}
+		return pem.EncodeToMemory(&pem.Block{
+			Bytes: publicBytes,
+			Type:  "PUBLIC KEY",
+		}), nil
 	}
 }
 
@@ -42,6 +65,28 @@ func DecodePrivateKey(private []byte, password string) (crypto.PrivateKey, error
 		return DecodeRsaPrivateKey(private, password)
 	case "EC PRIVATE KEY":
 		return DecodeEcdsaPrivateKey(private, password)
+	case "PRIVATE KEY":
+		var prikey crypto.PrivateKey
+		var err error
+		var bytes []byte
+		blk, _ := pem.Decode(private)
+		if blk == nil {
+			return nil, errors.New("private key error！")
+		}
+		if password != "" {
+			tmp, err := x509.DecryptPEMBlock(blk, []byte(password))
+			if err != nil {
+				return nil, err
+			}
+			bytes = tmp
+		} else {
+			bytes = blk.Bytes
+		}
+		prikey, err = x509.ParsePKCS8PrivateKey(bytes)
+		if err != nil {
+			return nil, err
+		}
+		return prikey, err
 	default:
 		return nil, errors.New("private key type error")
 	}
@@ -66,3 +111,7 @@ func EncodeSSHPublicKey(public crypto.PublicKey) ([]byte, error) {
 	}
 	return ssh.MarshalAuthorizedKey(publicKey), nil
 }
+
+func DecodeSSHPublicKey(pubStr []byte) (crypto.PublicKey, error) {
+	return ssh.ParsePublicKey(pubStr)
+}