Skip to content

Commit 3aceced

Browse files
authored
Add RSA encryption function to templates (#502)
* Add RSA encryption function to templates * fix the compile errors --------- Co-authored-by: rick <[email protected]>
1 parent 2e50d1c commit 3aceced

File tree

2 files changed

+82
-3
lines changed

2 files changed

+82
-3
lines changed

pkg/render/template.go

+35-3
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ You may obtain a copy of the License at
1010
Unless required by applicable law or agreed to in writing, software
1111
distributed under the License is distributed on an "AS IS" BASIS,
1212
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-
See the License for the specific language governing permissions and
13+
See the License for the specific language 24 permissions and
1414
limitations under the License.
1515
*/
1616
package render
@@ -19,16 +19,22 @@ import (
1919
"bytes"
2020
"context"
2121
"crypto/md5"
22+
"crypto/rand"
2223
"crypto/sha256"
2324
"encoding/base64"
2425
"encoding/hex"
2526
"encoding/json"
2627
"fmt"
2728
"html/template"
2829
"io"
29-
"math/rand"
30+
mathrand "math/rand"
3031
"strings"
3132

33+
"crypto/rsa"
34+
"crypto/x509"
35+
"encoding/pem"
36+
"errors"
37+
3238
"github.com/Masterminds/sprig/v3"
3339
"github.com/linuxsuren/api-testing/pkg/secret"
3440
"github.com/linuxsuren/api-testing/pkg/util"
@@ -85,6 +91,7 @@ func FuncMap() template.FuncMap {
8591
}
8692
funcs[item.FuncName] = item.Func
8793
}
94+
funcs["rasEncryptWithPublicKey"] = rasEncryptWithPublicKey
8895
return funcs
8996
}
9097

@@ -158,7 +165,7 @@ var advancedFuncs = []AdvancedFunc{{
158165
}, {
159166
FuncName: "randEnum",
160167
Func: func(items ...string) string {
161-
return items[rand.Intn(len(items))]
168+
return items[mathrand.Intn(len(items))]
162169
},
163170
}, {
164171
FuncName: "randEmail",
@@ -217,3 +224,28 @@ type AdvancedFunc struct {
217224
GoDogExper string
218225
Generator func(ctx context.Context, fields string) (err error)
219226
}
227+
228+
// rasEncryptWithPublicKey encrypts the given content with the provided public key
229+
func rasEncryptWithPublicKey(content, key string) (string, error) {
230+
block, _ := pem.Decode([]byte(key))
231+
if block == nil {
232+
return "", errors.New("failed to parse PEM block containing the public key")
233+
}
234+
235+
pub, err := x509.ParsePKIXPublicKey(block.Bytes)
236+
if err != nil {
237+
return "", fmt.Errorf("failed to parse DER encoded public key: %s", err)
238+
}
239+
240+
rsaPub, ok := pub.(*rsa.PublicKey)
241+
if !ok {
242+
return "", errors.New("key type is not RSA")
243+
}
244+
245+
encryptedData, err := rsa.EncryptPKCS1v15(rand.Reader, rsaPub, []byte(content))
246+
if err != nil {
247+
return "", fmt.Errorf("failed to encrypt with RSA public key: %s", err)
248+
}
249+
250+
return base64.StdEncoding.EncodeToString(encryptedData), nil
251+
}

pkg/render/template_test.go

+47
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ package render
1818
import (
1919
"bytes"
2020
"context"
21+
"crypto/rand"
22+
"crypto/rsa"
23+
"crypto/x509"
24+
"encoding/base64"
25+
"encoding/pem"
2126
"io"
2227
"testing"
2328

@@ -233,3 +238,45 @@ func TestSecret(t *testing.T) {
233238
assert.Equal(t, "hello", string(data))
234239
})
235240
}
241+
242+
func TestRasEncryptWithPublicKey(t *testing.T) {
243+
// Generate a new RSA key pair
244+
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
245+
if err != nil {
246+
t.Fatalf("Failed to generate private key: %v", err)
247+
}
248+
publicKey := &privateKey.PublicKey
249+
250+
// Encode the public key to PEM format
251+
pubASN1, err := x509.MarshalPKIXPublicKey(publicKey)
252+
if err != nil {
253+
t.Fatalf("Failed to marshal public key: %v", err)
254+
}
255+
pubBytes := pem.EncodeToMemory(&pem.Block{
256+
Type: "RSA PUBLIC KEY",
257+
Bytes: pubASN1,
258+
})
259+
260+
// Encrypt a message using the public key
261+
message := "hello world"
262+
encryptedMessage, err := rasEncryptWithPublicKey(message, string(pubBytes))
263+
if err != nil {
264+
t.Fatalf("Failed to encrypt message: %v", err)
265+
}
266+
267+
// Decrypt the message using the private key
268+
decodedMessage, err := base64.StdEncoding.DecodeString(encryptedMessage)
269+
if err != nil {
270+
t.Fatalf("Failed to decode message: %v", err)
271+
}
272+
decryptedBytes, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, decodedMessage)
273+
if err != nil {
274+
t.Fatalf("Failed to decrypt message: %v", err)
275+
}
276+
277+
// Verify the decrypted message
278+
decryptedMessage := string(decryptedBytes)
279+
if decryptedMessage != message {
280+
t.Fatalf("Decrypted message does not match original. Got: %s, want: %s", decryptedMessage, message)
281+
}
282+
}

0 commit comments

Comments
 (0)