arimelody.me/controller/totp.go

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 TOTP_TIME_STEP int64 = 30
const TOTP_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() / TOTP_TIME_STEP - int64(timeStepOffset)
    counterBytes := make([]byte, 8)
    binary.BigEndian.PutUint64(counterBytes, uint64(counter))

    mac := hmac.New(sha1.New, []byte(decodedSecret))
    mac.Write(counterBytes)
    hash := mac.Sum(nil)

    offset := hash[len(hash) - 1] & 0x0f
    binaryCode := int32(binary.BigEndian.Uint32(hash[offset : offset + 4]) & 0x7FFFFFFF)
    code := binaryCode % int32(math.Pow10(TOTP_CODE_LENGTH))

    return fmt.Sprintf(fmt.Sprintf("%%0%dd", TOTP_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",
        Host: "totp",
        Path: url.QueryEscape("arimelody.me") + ":" + url.QueryEscape(username),
    }

    query := url.Query()
    query.Set("secret", secret)
    query.Set("issuer", "arimelody.me")
    // query.Set("algorithm", "SHA1")
    // query.Set("digits", fmt.Sprintf("%d", TOTP_CODE_LENGTH))
    // query.Set("period", fmt.Sprintf("%d", TOTP_TIME_STEP))
    url.RawQuery = query.Encode()

    return url.String()
}

func GetTOTPsForAccount(db *sqlx.DB, accountID string) ([]model.TOTP, error) {
    totps := []model.TOTP{}

    err := db.Select(
        &totps,
        "SELECT * FROM totp " +
        "WHERE account=$1 " +
        "ORDER BY created_at ASC",
        accountID,
    )
    if err != nil {
        return nil, err
    }

    return totps, nil
}

func CheckTOTPForAccount(db *sqlx.DB, accountID string, totp string) (*model.TOTP, error) {
    totps, err := GetTOTPsForAccount(db, accountID)
    if err != nil {
        return nil, err
    }
    for _, method := range totps {
        check := GenerateTOTP(method.Secret, 0)
        if check == totp {
            return &method, nil
        }
        // try again with offset- maybe user input the code late?
        check = GenerateTOTP(method.Secret, 1)
        if check == totp {
            return &method, nil
        }
    }
    // user failed all TOTP checks
    // note: this state will still occur even if the account has no TOTP methods.
    return nil, nil
}

func GetTOTP(db *sqlx.DB, accountID string, name string) (*model.TOTP, error) {
    totp := model.TOTP{}

    err := db.Get(
        &totp,
        "SELECT * FROM totp " +
        "WHERE account=$1 AND name=$2",
        accountID,
        name,
    )
    if err != nil {
        if strings.Contains(err.Error(), "no rows") {
            return nil, nil
        }
        return nil, err
    }

    return &totp, nil
}

func CreateTOTP(db *sqlx.DB, totp *model.TOTP) error {
    _, err := db.Exec(
        "INSERT INTO totp (account, name, secret) " +
        "VALUES ($1,$2,$3)",
        totp.AccountID,
        totp.Name,
        totp.Secret,
    )
    return err
}

func DeleteTOTP(db *sqlx.DB, accountID string, name string) error {
    _, err := db.Exec(
        "DELETE FROM totp WHERE account=$1 AND name=$2",
        accountID,
        name,
    )
    return err
}