From 39b332b477e85d8688b8d58a50c6339a53f87fe6 Mon Sep 17 00:00:00 2001
From: ari melody <ari@arimelody.me>
Date: Tue, 21 Jan 2025 00:30:43 +0000
Subject: [PATCH] working TOTP codes YIPPEE

---
 controller/totp.go | 23 ++++++++++++++++++++++-
 main.go            |  2 +-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/controller/totp.go b/controller/totp.go
index 008634b..18616da 100644
--- a/controller/totp.go
+++ b/controller/totp.go
@@ -3,26 +3,35 @@ package controller
 import (
 	"arimelody-web/model"
 	"crypto/hmac"
+	"crypto/rand"
 	"crypto/sha1"
+	"encoding/base32"
 	"encoding/binary"
 	"fmt"
 	"math"
 	"net/url"
+	"os"
 	"strings"
 	"time"
 
 	"github.com/jmoiron/sqlx"
 )
 
+const TOTP_SECRET_LENGTH = 32
 const TIME_STEP int64 = 30
 const CODE_LENGTH = 6
 
 func GenerateTOTP(secret string, timeStepOffset int) string {
+    decodedSecret, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(secret)
+    if err != nil {
+        fmt.Fprintf(os.Stderr, "WARN: Invalid Base32 secret\n")
+    }
+
     counter := time.Now().Unix() / TIME_STEP - int64(timeStepOffset)
     counterBytes := make([]byte, 8)
     binary.BigEndian.PutUint64(counterBytes, uint64(counter))
 
-    mac := hmac.New(sha1.New, []byte(secret))
+    mac := hmac.New(sha1.New, []byte(decodedSecret))
     mac.Write(counterBytes)
     hash := mac.Sum(nil)
 
@@ -33,6 +42,18 @@ func GenerateTOTP(secret string, timeStepOffset int) string {
     return fmt.Sprintf(fmt.Sprintf("%%0%dd", CODE_LENGTH), code)
 }
 
+func GenerateTOTPSecret(length int) string {
+    bytes := make([]byte, length)
+    _, err := rand.Read(bytes)
+    if err != nil {
+        panic("FATAL: Failed to generate random TOTP bytes")
+    }
+
+    secret := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(bytes)
+
+    return strings.ToUpper(secret)
+}
+
 func GenerateTOTPURI(username string, secret string) string {
     url := url.URL{
         Scheme: "otpauth",
diff --git a/main.go b/main.go
index 7cf3363..e257da5 100644
--- a/main.go
+++ b/main.go
@@ -80,7 +80,7 @@ func main() {
             }
             username := os.Args[2]
             totpName := os.Args[3]
-            secret := controller.GenerateAlnumString(32)
+            secret := controller.GenerateTOTPSecret(controller.TOTP_SECRET_LENGTH)
 
             account, err := controller.GetAccount(global.DB, username)
             if err != nil {