package ssh

import (
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io"
	"net"
	"os"
)

func WithKey(username string, path string, keyCallBack ssh.HostKeyCallback) (ssh.ClientConfig, error) {
	signer, err := readKey(path)
	if err != nil {
		return ssh.ClientConfig{}, err
	}
	return AuthKey(username, []ssh.AuthMethod{ssh.PublicKeys(signer)}, keyCallBack)
}

func WithPassphrase(path string, passpharase []byte, username string, keyCallBack ssh.HostKeyCallback) (ssh.ClientConfig, error) {
	signer, err := readKeyWithPassphrase(path, passpharase)
	if err != nil {
		return ssh.ClientConfig{}, err
	}
	return AuthKey(username, []ssh.AuthMethod{ssh.PublicKeys(signer)}, keyCallBack)
}

func WithAgent(username string, keyCallBack ssh.HostKeyCallback) (ssh.ClientConfig, error) {
	socket := os.Getenv("SSH_AUTH_SOCK")
	conn, err := net.Dial("unix", socket)
	if err != nil {
		return ssh.ClientConfig{}, err
	}
	client := agent.NewClient(conn)
	return AuthKey(username, []ssh.AuthMethod{ssh.PublicKeysCallback(client.Signers)}, keyCallBack)
}

func WithPassword(username string, password string, keyCallBack ssh.HostKeyCallback) (ssh.ClientConfig, error) {
	return AuthKey(username, []ssh.AuthMethod{ssh.Password(password)}, keyCallBack)
}

func AuthKey(username string, auths []ssh.AuthMethod, keyCallBack ssh.HostKeyCallback) (ssh.ClientConfig, error) {
	return ssh.ClientConfig{
		User:            username,
		Auth:            auths,
		HostKeyCallback: keyCallBack,
	}, nil
}

func PublicKey(file string) ssh.AuthMethod {
	signer, err := readKey(file)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(signer)
}

func PublicKeys(filenames []string) []ssh.AuthMethod {
	var methods []ssh.AuthMethod
	for _, filename := range filenames {
		signer := PublicKey(filename)
		if signer != nil {
			methods = append(methods, signer)
		}
	}
	return methods
}

func readKey(filename string) (ssh.Signer, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	bytes, _ := io.ReadAll(f)
	return ssh.ParsePrivateKey(bytes)
}

func readKeyWithPassphrase(filename string, passpharase []byte) (ssh.Signer, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	bytes, _ := io.ReadAll(f)
	return ssh.ParsePrivateKeyWithPassphrase(bytes, passpharase)
}