From ece3a2a51fde420f5cdb34b254789957d612a924 Mon Sep 17 00:00:00 2001
From: Eashwar Ranganathan <eashwar@eashwar.com>
Date: Tue, 3 Dec 2024 21:37:56 -0800
Subject: [PATCH] Add support for Ed25519 keys

---
 go.mod   |  2 +-
 go.sum   |  4 ++--
 main.go  |  4 +++-
 setup.go | 32 +++++++++++++++++++++++++++++---
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/go.mod b/go.mod
index a4e62c7..3760ace 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module filippo.io/yubikey-agent
 go 1.19
 
 require (
-	github.com/go-piv/piv-go v1.10.0
+	github.com/go-piv/piv-go/v2 v2.3.0
 	github.com/twpayne/go-pinentry-minimal v0.0.0-20220113210447-2a5dc4396c2a
 	golang.org/x/crypto v0.4.0
 	golang.org/x/term v0.3.0
diff --git a/go.sum b/go.sum
index 465e662..fb5993c 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/go-piv/piv-go v1.10.0 h1:P1Y1VjBI5DnXW0+YkKmTuh5opWnMIrKriUaIOblee9Q=
-github.com/go-piv/piv-go v1.10.0/go.mod h1:NZ2zmjVkfFaL/CF8cVQ/pXdXtuj110zEKGdJM6fJZZM=
+github.com/go-piv/piv-go/v2 v2.3.0 h1:kKkrYlgLQTMPA6BiSL25A7/x4CEh2YCG7rtb/aTkx+g=
+github.com/go-piv/piv-go/v2 v2.3.0/go.mod h1:ShZi74nnrWNQEdWzRUd/3cSig3uNOcEZp+EWl0oewnI=
 github.com/twpayne/go-pinentry-minimal v0.0.0-20220113210447-2a5dc4396c2a h1:a1bRrtgkiv0tytmDVXU5Dqse/WOTws7JvsY2WxPMZ6M=
 github.com/twpayne/go-pinentry-minimal v0.0.0-20220113210447-2a5dc4396c2a/go.mod h1:ARJJXqNuaxVS84jX6ST52hQh0TtuQZWABhTe95a6BI4=
 golang.org/x/crypto v0.4.0 h1:UVQgzMY87xqpKNgb+kDsll2Igd33HszWHFLmpaRMq/8=
diff --git a/main.go b/main.go
index f50d239..ada75b3 100644
--- a/main.go
+++ b/main.go
@@ -10,6 +10,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/ecdsa"
+	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/rsa"
 	"errors"
@@ -28,7 +29,7 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/go-piv/piv-go/piv"
+	"github.com/go-piv/piv-go/v2/piv"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 	"golang.org/x/crypto/ssh/terminal"
@@ -249,6 +250,7 @@ func getPublicKey(yk *piv.YubiKey, slot piv.Slot) (ssh.PublicKey, error) {
 	}
 	switch cert.PublicKey.(type) {
 	case *ecdsa.PublicKey:
+	case ed25519.PublicKey:
 	case *rsa.PublicKey:
 	default:
 		return nil, fmt.Errorf("unexpected public key type: %T", cert.PublicKey)
diff --git a/setup.go b/setup.go
index ec6d10a..cdf866d 100644
--- a/setup.go
+++ b/setup.go
@@ -21,7 +21,7 @@ import (
 	"runtime/debug"
 	"time"
 
-	"github.com/go-piv/piv-go/piv"
+	"github.com/go-piv/piv-go/v2/piv"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/term"
 )
@@ -100,7 +100,15 @@ func runSetup(yk *piv.YubiKey) {
 	fmt.Println("")
 	fmt.Println("🧪 Reticulating splines...")
 
-	var key [24]byte
+	var version = yk.Version()
+	var key []byte
+	if supportsVersion(&version, 5, 4, 0) {
+		// Yubikey Firmware >=5.4.0 supports AES256 management keys
+		key = make([]byte, 32)
+	} else {
+		key = make([]byte, 24)
+	}
+
 	if _, err := rand.Read(key[:]); err != nil {
 		log.Fatal(err)
 	}
@@ -137,8 +145,16 @@ func runSetup(yk *piv.YubiKey) {
 		log.Fatalln("use --really-delete-all-piv-keys ⚠️")
 	}
 
+	var alg piv.Algorithm
+	if supportsVersion(&version, 5, 7, 0) {
+		// For newer Yubikeys, upgrade the key automatically to Ed25519
+		alg = piv.AlgorithmEd25519
+	} else {
+		alg = piv.AlgorithmEC256
+	}
+
 	pub, err := yk.GenerateKey(key, piv.SlotAuthentication, piv.Key{
-		Algorithm:   piv.AlgorithmEC256,
+		Algorithm:   alg,
 		PINPolicy:   piv.PINPolicyOnce,
 		TouchPolicy: piv.TouchPolicyAlways,
 	})
@@ -196,6 +212,16 @@ func runSetup(yk *piv.YubiKey) {
 	fmt.Println("💭 Remember: everything breaks, have a backup plan for when this YubiKey does.")
 }
 
+func supportsVersion(v *piv.Version, major, minor, patch int) bool {
+	if v.Major != major {
+		return v.Major > major
+	}
+	if v.Minor != minor {
+		return v.Minor > minor
+	}
+	return v.Patch >= patch
+}
+
 func randomSerialNumber() *big.Int {
 	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
 	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)