From 4811d506fa392aa60a5f5b6b66b4c84ab5e45ef4 Mon Sep 17 00:00:00 2001
From: Quantu <673216+Quantu@users.noreply.github.com>
Date: Tue, 9 Jul 2024 18:59:46 -0400
Subject: [PATCH] Add ability to change allowed number of PIN and PUK retries

---
 README.md          |  8 +++++++-
 v2/piv/piv.go      | 41 +++++++++++++++++++++++++++++++++++++++++
 v2/piv/piv_test.go | 15 +++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index b54e883..c912396 100644
--- a/README.md
+++ b/README.md
@@ -110,7 +110,13 @@ if _, err := io.ReadFull(rand.Reader, newKey); err != nil {
 newPIN := fmt.Sprintf("%06d", newPINInt)
 newPUK := fmt.Sprintf("%08d", newPUKInt)
 
-// Set all values to a new value.
+// If you want to change PIN/PUK retries, it's recommended to do it BEFORE changing
+// the PIN/PUK, as SetRetries will reset PIN/PUK to their default values.
+if err := yk.SetRetries(piv.DefaultManagementKey, piv.DefaultPIN, 5, 4); err != nil {
+	// ...
+}
+
+// Set all values to a new value. 
 if err := yk.SetManagementKey(piv.DefaultManagementKey, newKey); err != nil {
 	// ...
 }
diff --git a/v2/piv/piv.go b/v2/piv/piv.go
index f0ad16e..47669c2 100644
--- a/v2/piv/piv.go
+++ b/v2/piv/piv.go
@@ -664,6 +664,47 @@ func ykChangePUK(tx *scTx, oldPUK, newPUK string) error {
 	return err
 }
 
+// SetRetries sets the allowed retry count for the PIN and the PUK.
+//
+// Yubikeys allows one byte for storing each, allowed values are 1-255. In instances of greater
+// than 15 retries remaining, the remaining count will show 15 as Yubikeys only have 4 bits in
+// the response for remaining retries.
+//
+// IMPORTANT NOTE: Changing the retries on Yubikeys RESETS THE PIN AND PUK TO THEIR DEFAULTS!
+// If you use SetRetries, it is *highly* recommended that you follow it with SetPIN and SetPUK.
+// https://docs.yubico.com/yesdk/users-manual/application-piv/commands.html#set-pin-retries
+//
+//	if err := yk.SetRetries(piv.DefaultManagementKey, piv.DefaultPIN, 5, 4); err != nil {
+//		// ...
+//	}
+func (yk *YubiKey) SetRetries(managementKey []byte, pin string, pinRetries int, pukRetries int) error {
+	return ykSetRetries(yk.tx, managementKey, pin, pinRetries, pukRetries, yk.rand, yk.version)
+}
+
+func ykSetRetries(tx *scTx, managementKey []byte, pin string, pinRetries int, pukRetries int, rand io.Reader, version *version) error {
+	if pinRetries < 1 || pukRetries < 1 || pinRetries > 255 || pukRetries > 255 {
+		return fmt.Errorf("pinRetries and pukRetries must both be in range 1 - 255")
+	}
+
+	// NOTE: this action requires the management key AND PIN to be authenticated on
+	// the same transaction. It doesn't work otherwise.
+	if err := ykAuthenticate(tx, managementKey, rand, version); err != nil {
+		return fmt.Errorf("authenticating with management key: %w", err)
+	}
+	if err := ykLogin(tx, pin); err != nil {
+		return fmt.Errorf("authenticating with pin: %w", err)
+	}
+	cmd := apdu{
+		instruction: insSetPINRetries,
+		param1:      byte(pinRetries),
+		param2:      byte(pukRetries),
+	}
+	if _, err := tx.Transmit(cmd); err != nil {
+		return fmt.Errorf("command failed: %w", err)
+	}
+	return nil
+}
+
 func ykSelectApplication(tx *scTx, id []byte) error {
 	cmd := apdu{
 		instruction: insSelectApplication,
diff --git a/v2/piv/piv_test.go b/v2/piv/piv_test.go
index ee4ceed..02bfa33 100644
--- a/v2/piv/piv_test.go
+++ b/v2/piv/piv_test.go
@@ -281,6 +281,21 @@ func TestYubiKeyChangePUK(t *testing.T) {
 	}
 }
 
+func TestYubiKeyChangeRetries(t *testing.T) {
+	yk, close := newTestYubiKey(t)
+	defer close()
+
+	if err := yk.SetRetries(DefaultManagementKey, DefaultPIN, 3, 0); err == nil {
+		t.Errorf("successfully set retries to zeros, expected error")
+	}
+	if err := yk.SetRetries(DefaultManagementKey, DefaultPIN, 256, 3); err == nil {
+		t.Errorf("successfully set retries greater than 255, expected error")
+	}
+	if err := yk.SetRetries(DefaultManagementKey, DefaultPIN, 5, 3); err != nil {
+		t.Fatalf("setting pin/puk retries: %v", err)
+	}
+}
+
 func TestChangeManagementKey(t *testing.T) {
 	yk, close := newTestYubiKey(t)
 	defer close()