move-data-nas/internal/sshmanager/knownhosts.go

package sshmanager

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

type KnownHost struct {
	Host       string
	Port       int
	KeyType    string
	Fingerprint string
}

func NewKnownHostsCallback(knownHostsPath string, strictHostKeyChecking bool) (ssh.HostKeyCallback, error) {
	if !strictHostKeyChecking {
		return ssh.InsecureIgnoreHostKey(), nil
	}
	if knownHostsPath == "" {
		return ssh.InsecureIgnoreHostKey(), nil
	}
	_, err := os.Stat(knownHostsPath)
	if os.IsNotExist(err) {
		return ssh.InsecureIgnoreHostKey(), nil
	}
	if err != nil {
		return nil, err
	}
	return knownhosts.New(knownHostsPath)
}

func EnsureKnownHosts(sshDir string) (string, error) {
	path := filepath.Join(sshDir, "known_hosts")
	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDONLY, 0644)
	if err != nil {
		return "", err
	}
	f.Close()
	return path, nil
}

// AddKnownHost stores the host key in standard ssh known_hosts format (hostname keytype base64key).
// NOTE: Existing entries in known_hosts may need to be regenerated if they were stored in a different format.
func AddKnownHost(sshDir, host string, port int, keyData []byte) error {
	path := filepath.Join(sshDir, "known_hosts")
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	addr := host
	if port != 22 {
		addr = fmt.Sprintf("[%s]:%d", host, port)
	}

	pubKey, err := ssh.ParsePublicKey(keyData)
	if err != nil {
		return fmt.Errorf("parsing host key: %w", err)
	}
	line := fmt.Sprintf("%s %s\n", addr, strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pubKey))))
	if _, err := f.WriteString(line); err != nil {
		return err
	}
	return nil
}

func parseHostPort(entry string) (string, int) {
	if strings.HasPrefix(entry, "[") {
		var h string
		var p int
		if n, _ := fmt.Sscanf(entry, "[%[^]]]:%d", &h, &p); n == 2 {
			return h, p
		}
	}
	parts := strings.Split(entry, ":")
	if len(parts) == 2 {
		return parts[0], 22
	}
	return entry, 22
}

func GetKnownHost(sshDir, host string, port int) (*KnownHost, error) {
	path := filepath.Join(sshDir, "known_hosts")
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var targetAddr string
	if port != 22 {
		targetAddr = fmt.Sprintf("[%s]:%d", host, port)
	} else {
		targetAddr = host
	}

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) < 2 {
			continue
		}
		h, p := parseHostPort(parts[0])
		if (h == host || parts[0] == targetAddr) && p == port {
			return &KnownHost{
				Host:       h,
				Port:       p,
				KeyType:    parts[1],
				Fingerprint: parts[1] + " " + parts[2],
			}, nil
		}
	}
	return nil, nil
}

func HasKnownHost(sshDir, host string, port int) (bool, error) {
	kh, err := GetKnownHost(sshDir, host, port)
	if err != nil {
		return false, err
	}
	return kh != nil, nil
}

func RemoveKnownHost(sshDir, host string, port int) error {
	path := filepath.Join(sshDir, "known_hosts")
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	var lines []string
	targetAddr := host
	if port != 22 {
		targetAddr = fmt.Sprintf("[%s]:%d", host, port)
	}

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		h, p := parseHostPort(line)
		if h == host && p == port {
			continue
		}
		if line == targetAddr {
			continue
		}
		lines = append(lines, line)
	}

	tmp := path + ".tmp"
	wf, err := os.Create(tmp)
	if err != nil {
		return err
	}
	for _, l := range lines {
		wf.WriteString(l + "\n")
	}
	wf.Close()
	return os.Rename(tmp, path)
}