package imsvr

import (
	"github.com/xtaci/kcp-go"
	"net"
	"golang.org/x/crypto/ssh"
	"errors"
	"bytes"
)

const VERSION string = "1.0"

type Server struct{
	kcpListener *net.Listener
	sshConfig *ssh.ServerConfig
	isPasswdAuthSet bool
	cliAuthCallback func(username string, password []byte) bool
	sbnatAuthList map[string][]byte
}

func NewServer(conf *ServerConfig) (*Server, error) {
	kcplis,err := kcp.Listen(conf.BindAddress)
	if err != nil {
		return nil, err
	}
	sshcfg := &ssh.ServerConfig{
		NoClientAuth: false,
		MaxAuthTries: conf.MaxAuthTries,
		ServerVersion: "goSBNAT-Intermediate-Server-" + VERSION,
	}
	if conf.hostkeys == nil {
		return nil, errors.New("NoHostKeySpecified")
	}
	for _,v := range conf.hostkeys {
		sshcfg.AddHostKey(v)
	}
	svr := &Server{
		kcpListener: &kcplis,
		sshConfig: sshcfg,
	}
	if conf.ovrSSHKeyboardInteractiveCallback {
		svr.sshConfig.KeyboardInteractiveCallback = conf.ocSSHKeyboardInteractiveCallback
	}else {
		svr.sshConfig.KeyboardInteractiveCallback = svr.sshKbdIaCb
	}
	if conf.ovrSSHAuthLogCallback {
		svr.sshConfig.AuthLogCallback = conf.ocSSHAuthLogCallback
	}else{
		svr.sshConfig.AuthLogCallback = svr.sshAuthLogCb
	}
	if conf.ovrSSHPasswordCallback {
		svr.sshConfig.PasswordCallback = conf.ocSSHPasswordCallback
	}else {
		svr.sshConfig.PasswordCallback = svr.sshPwdCb
	}
	if conf.ovrSSHPublicKeyCallback {
		svr.sshConfig.PublicKeyCallback = conf.ocSSHPublicKeyCallback
	}else {
		svr.sshConfig.PublicKeyCallback = svr.sshPubKeyCb
	}
	svr.sbnatAuthList = make(map[string][]byte)
	return svr, nil
}

func (this *Server) sshKbdIaCb(conn ssh.ConnMetadata, client ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error){
	return nil, errors.New("KeyboardInteractiveAuthNotSupport")
}

func (this *Server) sshAuthLogCb(conn ssh.ConnMetadata, method string, err error){
	return
}

func (this *Server) sshPwdCb(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
	if this.cliAuthCallback == nil {
		return nil, errors.New("NoClientAuthCallback")
	}
	if this.cliAuthCallback(conn.User(), password) {
		perm := &ssh.Permissions{
			CriticalOptions: make(map[string]string),
			Extensions: make(map[string]string),
		}
		perm.CriticalOptions["type"] = "client"
		perm.CriticalOptions["user"] = conn.User()
		return perm, nil
	}else {
		return nil, errors.New("ClientAuthFailed")
	}
}

func (this *Server) sshPubKeyCb(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	v, ok := this.sbnatAuthList[conn.User()]
	if ok {
		if bytes.Equal(v, key.Marshal()) {
			perm := &ssh.Permissions{
				CriticalOptions: make(map[string]string),
				Extensions: make(map[string]string),
			}
			perm.CriticalOptions["type"] = "sbnat"
			perm.CriticalOptions["user"] = conn.User()
			return perm, nil
		}else {
			return nil, errors.New("FailedAuthSBNAT")
		}
	}else{
		return nil, errors.New("FailedAuthSBNAT")
	}
}

func (this *Server) AddAuthOfSBNAT(servername string, pubkey ssh.PublicKey) {
	this.sbnatAuthList[servername] = pubkey.Marshal()
}

func (this *Server) RegistClientAuthCallback(cb func(username string, password []byte) bool) {
	this.cliAuthCallback = cb
}