Merge pull request #42 from pquerna/pq/expose_setting_secret

Allow calls to {htop,totp}.Generate() to specify the secret
pquerna · Jun 1, 2019 · 43bebef · 43bebef
2 parents 1e101fa + ecc3285
commit 43bebef
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 10 deletions.
diff --git a/hotp/hotp.go b/hotp/hotp.go
@@ -146,6 +146,8 @@ type GenerateOpts struct {
 	AccountName string
 	// Size in size of the generated Secret. Defaults to 10 bytes.
 	SecretSize uint
+	// Secret to store. Defaults to a randomly generated secret of SecretSize.  You should generally leave this empty.
+	Secret []byte
 	// Digits to request. Defaults to 6.
 	Digits otp.Digits
 	// Algorithm to use for HMAC. Defaults to SHA1.
@@ -176,13 +178,17 @@ func Generate(opts GenerateOpts) (*otp.Key, error) {
 	// otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example
 
 	v := url.Values{}
-	secret := make([]byte, opts.SecretSize)
-	_, err := rand.Read(secret)
-	if err != nil {
-		return nil, err
+	if len(opts.Secret) != 0 {
+		v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
+	} else {
+		secret := make([]byte, opts.SecretSize)
+		_, err := rand.Read(secret)
+		if err != nil {
+			return nil, err
+		}
+		v.Set("secret", b32NoPadding.EncodeToString(secret))
 	}
 
-	v.Set("secret", b32NoPadding.EncodeToString(secret))
 	v.Set("issuer", opts.Issuer)
 	v.Set("algorithm", opts.Algorithm.String())
 	v.Set("digits", opts.Digits.String())

diff --git a/hotp/hotp_test.go b/hotp/hotp_test.go
@@ -169,4 +169,14 @@ func TestGenerate(t *testing.T) {
 	})
 	require.NoError(t, err, "Secret size is valid when length not divisable by 5.")
 	require.NotContains(t, k.Secret(), "=", "Secret has no escaped characters.")
+
+	k, err = Generate(GenerateOpts{
+		Issuer:      "SnakeOil",
+		AccountName: "alice@example.com",
+		Secret:      []byte("helloworld"),
+	})
+	require.NoError(t, err, "Secret generation failed")
+	sec, err := b32NoPadding.DecodeString(k.Secret())
+	require.NoError(t, err, "Secret wa not valid base32")
+	require.Equal(t, sec, []byte("helloworld"), "Specified Secret was not kept")
 }
diff --git a/totp/totp.go b/totp/totp.go
@@ -136,6 +136,8 @@ type GenerateOpts struct {
 	Period uint
 	// Size in size of the generated Secret. Defaults to 20 bytes.
 	SecretSize uint
+	// Secret to store. Defaults to a randomly generated secret of SecretSize.  You should generally leave this empty.
+	Secret []byte
 	// Digits to request. Defaults to 6.
 	Digits otp.Digits
 	// Algorithm to use for HMAC. Defaults to SHA1.
@@ -170,13 +172,17 @@ func Generate(opts GenerateOpts) (*otp.Key, error) {
 	// otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example
 
 	v := url.Values{}
-	secret := make([]byte, opts.SecretSize)
-	_, err := rand.Read(secret)
-	if err != nil {
-		return nil, err
+	if len(opts.Secret) != 0 {
+		v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
+	} else {
+		secret := make([]byte, opts.SecretSize)
+		_, err := rand.Read(secret)
+		if err != nil {
+			return nil, err
+		}
+		v.Set("secret", b32NoPadding.EncodeToString(secret))
 	}
 
-	v.Set("secret", b32NoPadding.EncodeToString(secret))
 	v.Set("issuer", opts.Issuer)
 	v.Set("period", strconv.FormatUint(uint64(opts.Period), 10))
 	v.Set("algorithm", opts.Algorithm.String())

diff --git a/totp/totp_test.go b/totp/totp_test.go
@@ -143,6 +143,16 @@ func TestGenerate(t *testing.T) {
 	})
 	require.NoError(t, err, "Secret size is valid when length not divisable by 5.")
 	require.NotContains(t, k.Secret(), "=", "Secret has no escaped characters.")
+
+	k, err = Generate(GenerateOpts{
+		Issuer:      "SnakeOil",
+		AccountName: "alice@example.com",
+		Secret:      []byte("helloworld"),
+	})
+	require.NoError(t, err, "Secret generation failed")
+	sec, err := b32NoPadding.DecodeString(k.Secret())
+	require.NoError(t, err, "Secret wa not valid base32")
+	require.Equal(t, sec, []byte("helloworld"), "Specified Secret was not kept")
 }
 
 func TestGoogleLowerCaseSecret(t *testing.T) {