Support custom hardware key prompt (#47273)

* Allow passing custom prompt to YubiKey

* Handle `prompt.Touch` cancellation

* Pass `HardwareKeyPrompt` through all the layers

* Add an empty `HardwareKeyPromptConstructor` to Connect

* Remove `ParsePrivateKeyWithCustomPrompt`

* Add missing godoc

* Fix teleterm tests

* Include `cliprompt.go` only for `go:build piv && !pivtest`

* Lint and test fixes

api/utils/keys/cliprompt.go

@@ -0,0 +1,126 @@
+//go:build piv && !pivtest
+
+// Copyright 2024 Gravitational, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package keys
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/go-piv/piv-go/piv"
+	"github.com/gravitational/trace"
+
+	"github.com/gravitational/teleport/api/utils/prompt"
+)
+
+type cliPrompt struct{}
+
+func (c *cliPrompt) AskPIN(ctx context.Context, message string) (string, error) {
+	password, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), message)
+	return password, trace.Wrap(err)
+}
+
+func (c *cliPrompt) Touch(_ context.Context) error {
+	_, err := fmt.Fprintln(os.Stderr, "Tap your YubiKey")
+	return trace.Wrap(err)
+}
+
+func (c *cliPrompt) ChangePIN(ctx context.Context) (*PINAndPUK, error) {
+	var pinAndPUK = &PINAndPUK{}
+	for {
+		fmt.Fprintf(os.Stderr, "Please set a new 6-8 character PIN.\n")
+		newPIN, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), "Enter your new YubiKey PIV PIN")
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		newPINConfirm, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), "Confirm your new YubiKey PIV PIN")
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+
+		if newPIN != newPINConfirm {
+			fmt.Fprintf(os.Stderr, "PINs do not match.\n")
+			continue
+		}
+
+		if newPIN == piv.DefaultPIN {
+			fmt.Fprintf(os.Stderr, "The default PIN %q is not supported.\n", piv.DefaultPIN)
+			continue
+		}
+
+		if !isPINLengthValid(newPIN) {
+			fmt.Fprintf(os.Stderr, "PIN must be 6-8 characters long.\n")
+			continue
+		}
+
+		pinAndPUK.PIN = newPIN
+		break
+	}
+
+	puk, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), "Enter your YubiKey PIV PUK to reset PIN [blank to use default PUK]")
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	pinAndPUK.PUK = puk
+
+	switch puk {
+	case piv.DefaultPUK:
+		fmt.Fprintf(os.Stderr, "The default PUK %q is not supported.\n", piv.DefaultPUK)
+		fallthrough
+	case "":
+		for {
+			fmt.Fprintf(os.Stderr, "Please set a new 6-8 character PUK (used to reset PIN).\n")
+			newPUK, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), "Enter your new YubiKey PIV PUK")
+			if err != nil {
+				return nil, trace.Wrap(err)
+			}
+			newPUKConfirm, err := prompt.Password(ctx, os.Stderr, prompt.Stdin(), "Confirm your new YubiKey PIV PUK")
+			if err != nil {
+				return nil, trace.Wrap(err)
+			}
+
+			if newPUK != newPUKConfirm {
+				fmt.Fprintf(os.Stderr, "PUKs do not match.\n")
+				continue
+			}
+
+			if newPUK == piv.DefaultPUK {
+				fmt.Fprintf(os.Stderr, "The default PUK %q is not supported.\n", piv.DefaultPUK)
+				continue
+			}
+
+			if !isPINLengthValid(newPUK) {
+				fmt.Fprintf(os.Stderr, "PUK must be 6-8 characters long.\n")
+				continue
+			}
+
+			pinAndPUK.PUK = newPUK
+			pinAndPUK.PUKChanged = true
+			break
+		}
+	}
+	return pinAndPUK, nil
+}
+
+func (c *cliPrompt) ConfirmSlotOverwrite(ctx context.Context, message string) (bool, error) {
+	confirmation, err := prompt.Confirmation(ctx, os.Stderr, prompt.Stdin(), message)
+	return confirmation, trace.Wrap(err)
+}
+
+func isPINLengthValid(pin string) bool {
+	return len(pin) >= 6 && len(pin) <= 8
+}

api/utils/keys/privatekey.go

@@ -212,16 +212,40 @@ func LoadPrivateKey(keyFile string) (*PrivateKey, error) {
 	return priv, nil
 }
 
+// ParsePrivateKeyOptions contains config options for ParsePrivateKey.
+type ParsePrivateKeyOptions struct {
+	// CustomHardwareKeyPrompt is a custom hardware key prompt to use when asking
+	// for a hardware key PIN, touch, etc.
+	// If empty, a default CLI prompt is used.
+	CustomHardwareKeyPrompt HardwareKeyPrompt
+}
+
+// ParsePrivateKeyOpt applies configuration options.
+type ParsePrivateKeyOpt func(o *ParsePrivateKeyOptions)
+
+// WithCustomPrompt sets a custom hardware key prompt.
+func WithCustomPrompt(prompt HardwareKeyPrompt) ParsePrivateKeyOpt {
+	return func(o *ParsePrivateKeyOptions) {
+		o.CustomHardwareKeyPrompt = prompt
+	}
+}
+
 // ParsePrivateKey returns the PrivateKey for the given key PEM block.
-func ParsePrivateKey(keyPEM []byte) (*PrivateKey, error) {
+// Allows passing a custom hardware key prompt.
+func ParsePrivateKey(keyPEM []byte, opts ...ParsePrivateKeyOpt) (*PrivateKey, error) {
+	var appliedOpts ParsePrivateKeyOptions
+	for _, o := range opts {
+		o(&appliedOpts)
+	}
+
 	block, _ := pem.Decode(keyPEM)
 	if block == nil {
 		return nil, trace.BadParameter("expected PEM encoded private key")
 	}
 
 	switch block.Type {
 	case pivYubiKeyPrivateKeyType:
-		priv, err := parseYubiKeyPrivateKeyData(block.Bytes)
+		priv, err := parseYubiKeyPrivateKeyData(block.Bytes, appliedOpts.CustomHardwareKeyPrompt)
 		return priv, trace.Wrap(err, "parsing YubiKey private key")
 	case OpenSSHPrivateKeyType:
 		priv, err := ssh.ParseRawPrivateKey(keyPEM)
@@ -299,7 +323,7 @@ func MarshalPrivateKey(key crypto.Signer) ([]byte, error) {
 }
 
 // LoadKeyPair returns the PrivateKey for the given private and public key files.
-func LoadKeyPair(privFile, sshPubFile string) (*PrivateKey, error) {
+func LoadKeyPair(privFile, sshPubFile string, customPrompt HardwareKeyPrompt) (*PrivateKey, error) {
 	privPEM, err := os.ReadFile(privFile)
 	if err != nil {
 		return nil, trace.ConvertSystemError(err)
@@ -310,16 +334,16 @@ func LoadKeyPair(privFile, sshPubFile string) (*PrivateKey, error) {
 		return nil, trace.ConvertSystemError(err)
 	}
 
-	priv, err := ParseKeyPair(privPEM, marshaledSSHPub)
+	priv, err := ParseKeyPair(privPEM, marshaledSSHPub, customPrompt)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	return priv, nil
 }
 
 // ParseKeyPair returns the PrivateKey for the given private and public key PEM blocks.
-func ParseKeyPair(privPEM, marshaledSSHPub []byte) (*PrivateKey, error) {
-	priv, err := ParsePrivateKey(privPEM)
+func ParseKeyPair(privPEM, marshaledSSHPub []byte, customPrompt HardwareKeyPrompt) (*PrivateKey, error) {
+	priv, err := ParsePrivateKey(privPEM, WithCustomPrompt(customPrompt))
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}