From 016b032d0a54d13a4fbda4a72ac2d435e376348e Mon Sep 17 00:00:00 2001
From: Ke Zhu <davidzhu.4.2@gmail.com>
Date: Thu, 13 Aug 2015 22:37:29 -0400
Subject: [PATCH] RSA to RSA-OAEP

---
 pkg/utils/sshutil/sshutil.go      | 24 +++++++++++++++-------
 pkg/utils/sshutil/sshutil_test.go | 33 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 7 deletions(-)
 create mode 100644 pkg/utils/sshutil/sshutil_test.go

diff --git a/pkg/utils/sshutil/sshutil.go b/pkg/utils/sshutil/sshutil.go
index 197265f00..1e397b800 100644
--- a/pkg/utils/sshutil/sshutil.go
+++ b/pkg/utils/sshutil/sshutil.go
@@ -4,7 +4,9 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
+	"encoding/base64"
 	"encoding/pem"
+	"hash"
 
 	"github.com/drone/drone/Godeps/_workspace/src/code.google.com/p/go.crypto/ssh"
 )
@@ -38,15 +40,23 @@ func MarshalPrivateKey(privkey *rsa.PrivateKey) []byte {
 	return privateKeyPEM
 }
 
-// helper function to encrypt a plain-text string using
+// Encrypt is helper function to encrypt a plain-text string using
 // an RSA public key.
-func Encrypt(pubkey *rsa.PublicKey, msg string) ([]byte, error) {
-	return rsa.EncryptPKCS1v15(rand.Reader, pubkey, []byte(msg))
+func Encrypt(hash hash.Hash, pubkey *rsa.PublicKey, msg string) (string, error) {
+	src, err := rsa.EncryptOAEP(hash, rand.Reader, pubkey, []byte(msg), nil)
+
+	return base64.StdEncoding.EncodeToString(src), err
 }
 
-// helper function to encrypt a plain-text string using
+// Decrypt is helper function to encrypt a plain-text string using
 // an RSA public key.
-func Decrypt(privkey *rsa.PrivateKey, secret string) (string, error) {
-	msg, err := rsa.DecryptPKCS1v15(rand.Reader, privkey, []byte(secret))
-	return string(msg), err
+func Decrypt(hash hash.Hash, privkey *rsa.PrivateKey, secret string) (string, error) {
+	decoded, err := base64.StdEncoding.DecodeString(secret)
+	if err != nil {
+		return "", err
+	}
+
+	out, err := rsa.DecryptOAEP(hash, rand.Reader, privkey, decoded, nil)
+
+	return string(out), err
 }
diff --git a/pkg/utils/sshutil/sshutil_test.go b/pkg/utils/sshutil/sshutil_test.go
new file mode 100644
index 000000000..7405de5d6
--- /dev/null
+++ b/pkg/utils/sshutil/sshutil_test.go
@@ -0,0 +1,33 @@
+package sshutil
+
+import (
+	"crypto/sha256"
+	"testing"
+
+	"github.com/drone/drone/Godeps/_workspace/src/github.com/franela/goblin"
+)
+
+func TestSSHUtil(t *testing.T) {
+
+	g := goblin.Goblin(t)
+	g.Describe("sshutil", func() {
+		var encrypted, testMsg string
+
+		privkey, err := GeneratePrivateKey()
+		g.Assert(err == nil).IsTrue()
+		pubkey := privkey.PublicKey
+		sha256 := sha256.New()
+		testMsg = "foo=bar"
+
+		g.Before(func() {
+			encrypted, err = Encrypt(sha256, &pubkey, testMsg)
+			g.Assert(err == nil).IsTrue()
+		})
+
+		g.It("Can decrypt encrypted msg", func() {
+			decrypted, err := Decrypt(sha256, privkey, encrypted)
+			g.Assert(err == nil).IsTrue()
+			g.Assert(decrypted == testMsg).IsTrue()
+		})
+	})
+}