1package sshutils
 2
 3import (
 4	"bytes"
 5	"context"
 6	"fmt"
 7
 8	"github.com/charmbracelet/ssh"
 9	gossh "golang.org/x/crypto/ssh"
10)
11
12// ParseAuthorizedKey parses an authorized key string into a public key.
13func ParseAuthorizedKey(ak string) (gossh.PublicKey, string, error) {
14	pk, c, _, _, err := gossh.ParseAuthorizedKey([]byte(ak))
15	if err != nil {
16		return nil, "", fmt.Errorf("failed to parse authorized key: %w", err)
17	}
18	return pk, c, nil
19}
20
21// MarshalAuthorizedKey marshals a public key into an authorized key string.
22//
23// This is the inverse of ParseAuthorizedKey.
24// This function is a copy of ssh.MarshalAuthorizedKey, but without the trailing newline.
25// It returns an empty string if pk is nil.
26func MarshalAuthorizedKey(pk gossh.PublicKey) string {
27	if pk == nil {
28		return ""
29	}
30	return string(bytes.TrimSuffix(gossh.MarshalAuthorizedKey(pk), []byte("\n")))
31}
32
33// KeysEqual returns whether the two public keys are equal.
34func KeysEqual(a, b gossh.PublicKey) bool {
35	return ssh.KeysEqual(a, b)
36}
37
38// PublicKeyFromContext returns the public key from the context.
39func PublicKeyFromContext(ctx context.Context) gossh.PublicKey {
40	if pk, ok := ctx.Value(ssh.ContextKeyPublicKey).(gossh.PublicKey); ok {
41		return pk
42	}
43	return nil
44}
45
46// ContextKeySession is the context key for the SSH session.
47var ContextKeySession = &struct{ string }{"session"}
48
49// SessionFromContext returns the SSH session from the context.
50func SessionFromContext(ctx context.Context) ssh.Session {
51	if s, ok := ctx.Value(ContextKeySession).(ssh.Session); ok {
52		return s
53	}
54	return nil
55}