1// Package sshutils provides SSH utility functions.
 2package sshutils
 3
 4import (
 5	"bytes"
 6	"context"
 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	return pk, c, err //nolint:wrapcheck
16}
17
18// MarshalAuthorizedKey marshals a public key into an authorized key string.
19//
20// This is the inverse of ParseAuthorizedKey.
21// This function is a copy of ssh.MarshalAuthorizedKey, but without the trailing newline.
22// It returns an empty string if pk is nil.
23func MarshalAuthorizedKey(pk gossh.PublicKey) string {
24	if pk == nil {
25		return ""
26	}
27	return string(bytes.TrimSuffix(gossh.MarshalAuthorizedKey(pk), []byte("\n")))
28}
29
30// KeysEqual returns whether the two public keys are equal.
31func KeysEqual(a, b gossh.PublicKey) bool {
32	return ssh.KeysEqual(a, b)
33}
34
35// PublicKeyFromContext returns the public key from the context.
36func PublicKeyFromContext(ctx context.Context) gossh.PublicKey {
37	if pk, ok := ctx.Value(ssh.ContextKeyPublicKey).(gossh.PublicKey); ok {
38		return pk
39	}
40	return nil
41}
42
43// ContextKeySession is the context key for the SSH session.
44var ContextKeySession = &struct{ string }{"session"}
45
46// SessionFromContext returns the SSH session from the context.
47func SessionFromContext(ctx context.Context) ssh.Session {
48	if s, ok := ctx.Value(ContextKeySession).(ssh.Session); ok {
49		return s
50	}
51	return nil
52}