package sshtun

import (
	"bytes"
	"fmt"
	"net"

	"github.com/Neur0toxine/sshpoke/pkg/proto/ssh"
)

func FixedHostKeys(keys []ssh.PublicKey) ssh.HostKeyCallback {
	m := make(map[string]ssh.PublicKey)
	for _, key := range keys {
		m[key.Type()] = key
	}
	hk := &fixedHostKeys{keys: m}
	return hk.check
}

type fixedHostKeys struct {
	keys map[string]ssh.PublicKey
}

func (f *fixedHostKeys) check(hostname string, remote net.Addr, key ssh.PublicKey) error {
	if f.keys == nil {
		return fmt.Errorf("ssh: host keys should be defined")
	}
	if len(f.keys) == 0 {
		return fmt.Errorf("ssh: no host keys were provided")
	}
	hostKey, found := f.keys[key.Type()]
	if !found || !bytes.Equal(key.Marshal(), hostKey.Marshal()) {
		return fmt.Errorf("ssh: host key mismatch")
	}
	return nil
}

func CombineHostKeyCallbacks(callbacks ...ssh.HostKeyCallback) ssh.HostKeyCallback {
	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		var err error
		for _, cb := range callbacks {
			err = cb(hostname, remote, key)
			if err == nil {
				return nil
			}
		}
		return err
	}
}