From 31823eea3a23b8e2da60a6212ad78526d973d29a Mon Sep 17 00:00:00 2001
From: Morten Linderud <morten@linderud.pw>
Date: Tue, 22 Aug 2023 22:21:08 +0200
Subject: [PATCH] key: Implement Fingerprint and AuthorizedKey

Signed-off-by: Morten Linderud <morten@linderud.pw>
---
 agent/agent.go             | 19 +++----------------
 cmd/ssh-tpm-keygen/main.go | 15 +++------------
 key/key.go                 | 24 ++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 6c7cbee..88712ea 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -61,12 +61,7 @@ func (a *Agent) AddTPMKey(contents []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	sshpubkey, err := k.SSHPublicKey()
-	if err != nil {
-		return nil, err
-	}
-
-	a.keys[ssh.FingerprintSHA256(sshpubkey)] = k
+	a.keys[k.Fingerprint()] = k
 
 	return []byte(""), nil
 }
@@ -230,11 +225,7 @@ func (a *Agent) serve() {
 
 func (a *Agent) AddKey(k *key.Key) error {
 	slog.Debug("called addkey")
-	sshpubkey, err := k.SSHPublicKey()
-	if err != nil {
-		return err
-	}
-	a.keys[ssh.FingerprintSHA256(sshpubkey)] = k
+	a.keys[k.Fingerprint()] = k
 	return nil
 }
 
@@ -296,11 +287,7 @@ func LoadKeys(keyDir string) (map[string]*key.Key, error) {
 				slog.Debug("%s not a TPM sealed key: %v\n", path, err)
 				return nil
 			}
-			sshpubkey, err := k.SSHPublicKey()
-			if err != nil {
-				return fmt.Errorf("%s can't read ssh public key from TPM public: %v", path, err)
-			}
-			keys[ssh.FingerprintSHA256(sshpubkey)] = k
+			keys[k.Fingerprint()] = k
 			return nil
 		},
 	)
diff --git a/cmd/ssh-tpm-keygen/main.go b/cmd/ssh-tpm-keygen/main.go
index 1a8b841..dcefb20 100644
--- a/cmd/ssh-tpm-keygen/main.go
+++ b/cmd/ssh-tpm-keygen/main.go
@@ -275,22 +275,13 @@ func main() {
 		}
 	}
 
-	sshKey, err = k.SSHPublicKey()
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	pubkeyLine :=
-		strings.TrimSuffix(string(ssh.MarshalAuthorizedKey(sshKey)), "\n") +
-			" " + comment + "\n"
-
 	if importKey == "" {
-		if err := os.WriteFile(pubkeyFilename, []byte(pubkeyLine), 0644); err != nil {
+		if err := os.WriteFile(pubkeyFilename, k.AuthorizedKey(), 0600); err != nil {
 			log.Fatal(err)
 		}
 	}
 
-	if err := os.WriteFile(privatekeyFilename, key.EncodeKey(k), 0600); err != nil {
+	if err := os.WriteFile(privatekeyFilename, k.Encode(), 0600); err != nil {
 		log.Fatal(err)
 	}
 
@@ -299,6 +290,6 @@ func main() {
 		fmt.Printf("Your public key has been saved in %s\n", pubkeyFilename)
 	}
 	fmt.Printf("The key fingerprint is:\n")
-	fmt.Println(ssh.FingerprintSHA256(sshKey))
+	fmt.Println(k.Fingerprint())
 	fmt.Println("The key's randomart image is the color of television, tuned to a dead channel.")
 }
diff --git a/key/key.go b/key/key.go
index e9fe726..1171d0f 100644
--- a/key/key.go
+++ b/key/key.go
@@ -9,6 +9,7 @@ import (
 	"encoding/pem"
 	"fmt"
 	"math/big"
+	"strings"
 
 	"github.com/foxboron/ssh-tpm-agent/utils"
 	"github.com/google/go-tpm/tpm2"
@@ -97,6 +98,29 @@ func (k *Key) SSHPublicKey() (ssh.PublicKey, error) {
 	return ssh.NewPublicKey(pubkey)
 }
 
+func (k *Key) Fingerprint() string {
+	sshKey, err := k.SSHPublicKey()
+	if err != nil {
+		// This shouldn't happen
+		panic("not a valid ssh key")
+	}
+	return ssh.FingerprintSHA256(sshKey)
+}
+
+func (k *Key) AuthorizedKey() []byte {
+	sshKey, err := k.SSHPublicKey()
+	if err != nil {
+		// This shouldn't happen
+		panic("not a valid ssh key")
+	}
+	authKey := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshKey)))
+	return []byte(fmt.Sprintf("%s %s\n", authKey, string(k.Comment)))
+}
+
+func (k *Key) Encode() []byte {
+	return EncodeKey(k)
+}
+
 func UnmarshalKey(b []byte) (*Key, error) {
 	var key Key
 	var comment []byte