Restrict echo to only allow characters 0-9

kittens/ssh/main.go

@@ -562,7 +562,7 @@ func drain_potential_tty_garbage(term *tty.Term) {
 	if err != nil {
 		return
 	}
-	canary, err := secrets.TokenHex()
+	canary, err := secrets.TokenAlphabet(secrets.DEFAULT_NUM_OF_BYTES_FOR_TOKEN, "0123456789")
 	if err != nil {
 		return
 	}

tools/utils/secrets/tokens.go

@@ -6,6 +6,7 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
+	"math/big"
 
 	"github.com/emmansun/base64"
 )
@@ -14,6 +15,39 @@ var _ = fmt.Print
 
 const DEFAULT_NUM_OF_BYTES_FOR_TOKEN = 32
 
+func encode_bytes(input []byte, alphabet string) string {
+	if len(input) == 0 {
+		return ""
+	}
+
+	base := big.NewInt(int64(len(alphabet)))
+	num := new(big.Int).SetBytes(input)
+	mod := new(big.Int)
+
+	var result []byte
+
+	// Convert the large integer into the custom base system
+	for num.Sign() > 0 {
+		num.DivMod(num, base, mod)
+		result = append(result, alphabet[mod.Int64()])
+	}
+
+	// Preserve leading zeros from the original byte slice
+	for _, b := range input {
+		if b != 0x00 {
+			break
+		}
+		result = append(result, alphabet[0])
+	}
+
+	// Reverse the result slice since DivMod extracts digits backward
+	for i, j := 0, len(result)-1; i < j; i, j = i+1, j-1 {
+		result[i], result[j] = result[j], result[i]
+	}
+
+	return string(result)
+}
+
 func TokenBytes(nbytes ...int) ([]byte, error) {
 	if len(nbytes) == 0 {
 		nbytes = []int{DEFAULT_NUM_OF_BYTES_FOR_TOKEN}
@@ -41,3 +75,11 @@ func TokenBase64(nbytes ...int) (string, error) {
 	}
 	return base64.StdEncoding.EncodeToString(b), nil
 }
+
+func TokenAlphabet(nbytes int, alphabet string) (string, error) {
+	b, err := TokenBytes(nbytes)
+	if err != nil {
+		return "", err
+	}
+	return encode_bytes(b, alphabet), nil
+}