From 2795260d2189d89c843c918d009b4a1d9472e6ea Mon Sep 17 00:00:00 2001
From: meshmayhem <meshmayhem@airmail.cc>
Date: Wed, 18 Sep 2024 16:01:48 +0000
Subject: [PATCH] Improve genkeys with bit selection

---
 cmd/genkeys/main.go  | 48 +++++++++++++++++++++++++++++---------------
 src/config/config.go | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 16 deletions(-)

diff --git a/cmd/genkeys/main.go b/cmd/genkeys/main.go
index 36107c0a..d3b76193 100644
--- a/cmd/genkeys/main.go
+++ b/cmd/genkeys/main.go
@@ -13,12 +13,14 @@ package main
 import (
 	"crypto/ed25519"
 	"encoding/hex"
+	"flag"
 	"fmt"
 	"net"
 	"runtime"
 	"time"
 
 	"github.com/yggdrasil-network/yggdrasil-go/src/address"
+	"github.com/yggdrasil-network/yggdrasil-go/src/config"
 )
 
 type keySet struct {
@@ -27,23 +29,37 @@ type keySet struct {
 }
 
 func main() {
-	threads := runtime.GOMAXPROCS(0)
-	fmt.Println("Threads:", threads)
+	security := flag.Int("security", 0, "generates a key with a specific amount of security bits. defaults to 0 which continuously generates more keys")
+	flag.Parse()
+
 	start := time.Now()
-	var currentBest ed25519.PublicKey
-	newKeys := make(chan keySet, threads)
-	for i := 0; i < threads; i++ {
-		go doKeys(newKeys)
-	}
-	for {
-		newKey := <-newKeys
-		if isBetter(currentBest, newKey.pub) || len(currentBest) == 0 {
-			currentBest = newKey.pub
-			fmt.Println("-----", time.Since(start))
-			fmt.Println("Priv:", hex.EncodeToString(newKey.priv))
-			fmt.Println("Pub:", hex.EncodeToString(newKey.pub))
-			addr := address.AddrForKey(newKey.pub)
-			fmt.Println("IP:", net.IP(addr[:]).String())
+	if (*security > 0) {
+		// If higher than 0, generates a key with the set amount of security bits
+		var secureKey keySet
+		secureKey.priv, secureKey.pub = config.NewSecureKeyPair(*security)
+		fmt.Println("-----", time.Since(start))
+		fmt.Println("Priv:", hex.EncodeToString(secureKey.priv))
+		fmt.Println("Pub:", hex.EncodeToString(secureKey.pub))
+		addr := address.AddrForKey(secureKey.pub)
+		fmt.Println("IP:", net.IP(addr[:]).String())
+	} else {
+		threads := runtime.GOMAXPROCS(0)
+		fmt.Println("Threads:", threads)
+		var currentBest ed25519.PublicKey
+		newKeys := make(chan keySet, threads)
+		for i := 0; i < threads; i++ {
+			go doKeys(newKeys)
+		}
+		for {
+			newKey := <-newKeys
+			if isBetter(currentBest, newKey.pub) || len(currentBest) == 0 {
+				currentBest = newKey.pub
+				fmt.Println("-----", time.Since(start))
+				fmt.Println("Priv:", hex.EncodeToString(newKey.priv))
+				fmt.Println("Pub:", hex.EncodeToString(newKey.pub))
+				addr := address.AddrForKey(newKey.pub)
+				fmt.Println("IP:", net.IP(addr[:]).String())
+			}
 		}
 	}
 }
diff --git a/src/config/config.go b/src/config/config.go
index 9a7f7180..54d31f27 100644
--- a/src/config/config.go
+++ b/src/config/config.go
@@ -30,6 +30,8 @@ import (
 	"io"
 	"math/big"
 	"os"
+	"runtime"
+	"strings"
 	"time"
 
 	"github.com/hjson/hjson-go/v4"
@@ -208,6 +210,41 @@ func (cfg *NodeConfig) NewPrivateKey() {
 	cfg.PrivateKey = KeyBytes(spriv)
 }
 
+func NewSecureKeyPair(bits int) (priv ed25519.PrivateKey, pub ed25519.PublicKey) {
+	// Generates a key pair with a prescribed number of security bits.
+	threads := runtime.GOMAXPROCS(0)
+	if (bits > 64) {
+		bits = 64
+		// Bounding the maximum number of security bits to the maximum public key length of 64.
+	}
+	type keySet struct {
+		priv ed25519.PrivateKey
+		pub  ed25519.PublicKey
+	}
+	expected :=  strings.Repeat("0", bits)
+	// Generates the expected security substring in advance
+	newKeys := make(chan keySet, threads)
+	for i := 0; i < threads; i++ {
+		go func(out chan<- keySet) {
+			for {
+				pub, priv, err := ed25519.GenerateKey(nil)
+				if err != nil {
+					panic(err)
+				}
+				if !(hex.EncodeToString(pub)[0:bits] == expected) {
+					// Checks if the public key contains the expected security substring
+					continue
+				}
+				out <- keySet{priv, pub}
+			}
+		}(newKeys)
+	}
+	for {
+		newKey := <-newKeys
+		return newKey.priv, newKey.pub
+	}
+}
+
 func (cfg *NodeConfig) MarshalPEMPrivateKey() ([]byte, error) {
 	b, err := x509.MarshalPKCS8PrivateKey(ed25519.PrivateKey(cfg.PrivateKey))
 	if err != nil {