Skip to content

Commit

Permalink
Add RSA encryption function to templates (#502)
Browse files Browse the repository at this point in the history
* Add RSA encryption function to templates

* fix the compile errors

---------

Co-authored-by: rick <LinuxSuRen@users.noreply.github.com>
  • Loading branch information
LinuxSuRen and LinuxSuRen committed Jul 3, 2024
1 parent 2e50d1c commit 3aceced
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 3 deletions.
38 changes: 35 additions & 3 deletions pkg/render/template.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ You may obtain a copy of the License at
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
See the License for the specific language 24 permissions and
limitations under the License.
*/
package render
Expand All @@ -19,16 +19,22 @@ import (
"bytes"
"context"
"crypto/md5"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"html/template"
"io"
"math/rand"
mathrand "math/rand"
"strings"

"crypto/rsa"
"crypto/x509"
"encoding/pem"
"errors"

"github.com/Masterminds/sprig/v3"
"github.com/linuxsuren/api-testing/pkg/secret"
"github.com/linuxsuren/api-testing/pkg/util"
Expand Down Expand Up @@ -85,6 +91,7 @@ func FuncMap() template.FuncMap {
}
funcs[item.FuncName] = item.Func
}
funcs["rasEncryptWithPublicKey"] = rasEncryptWithPublicKey
return funcs
}

Expand Down Expand Up @@ -158,7 +165,7 @@ var advancedFuncs = []AdvancedFunc{{
}, {
FuncName: "randEnum",
Func: func(items ...string) string {
return items[rand.Intn(len(items))]
return items[mathrand.Intn(len(items))]
},
}, {
FuncName: "randEmail",
Expand Down Expand Up @@ -217,3 +224,28 @@ type AdvancedFunc struct {
GoDogExper string
Generator func(ctx context.Context, fields string) (err error)
}

// rasEncryptWithPublicKey encrypts the given content with the provided public key
func rasEncryptWithPublicKey(content, key string) (string, error) {
block, _ := pem.Decode([]byte(key))
if block == nil {
return "", errors.New("failed to parse PEM block containing the public key")
}

pub, err := x509.ParsePKIXPublicKey(block.Bytes)
if err != nil {
return "", fmt.Errorf("failed to parse DER encoded public key: %s", err)
}

rsaPub, ok := pub.(*rsa.PublicKey)
if !ok {
return "", errors.New("key type is not RSA")
}

encryptedData, err := rsa.EncryptPKCS1v15(rand.Reader, rsaPub, []byte(content))
if err != nil {
return "", fmt.Errorf("failed to encrypt with RSA public key: %s", err)
}

return base64.StdEncoding.EncodeToString(encryptedData), nil
}
47 changes: 47 additions & 0 deletions pkg/render/template_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@ package render
import (
"bytes"
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/base64"
"encoding/pem"
"io"
"testing"

Expand Down Expand Up @@ -233,3 +238,45 @@ func TestSecret(t *testing.T) {
assert.Equal(t, "hello", string(data))
})
}

func TestRasEncryptWithPublicKey(t *testing.T) {
// Generate a new RSA key pair
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("Failed to generate private key: %v", err)
}
publicKey := &privateKey.PublicKey

// Encode the public key to PEM format
pubASN1, err := x509.MarshalPKIXPublicKey(publicKey)
if err != nil {
t.Fatalf("Failed to marshal public key: %v", err)
}
pubBytes := pem.EncodeToMemory(&pem.Block{
Type: "RSA PUBLIC KEY",
Bytes: pubASN1,
})

// Encrypt a message using the public key
message := "hello world"
encryptedMessage, err := rasEncryptWithPublicKey(message, string(pubBytes))
if err != nil {
t.Fatalf("Failed to encrypt message: %v", err)
}

// Decrypt the message using the private key
decodedMessage, err := base64.StdEncoding.DecodeString(encryptedMessage)
if err != nil {
t.Fatalf("Failed to decode message: %v", err)
}
decryptedBytes, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, decodedMessage)
if err != nil {
t.Fatalf("Failed to decrypt message: %v", err)
}

// Verify the decrypted message
decryptedMessage := string(decryptedBytes)
if decryptedMessage != message {
t.Fatalf("Decrypted message does not match original. Got: %s, want: %s", decryptedMessage, message)
}
}

0 comments on commit 3aceced

Please sign in to comment.