package zsshrpc_server

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"os"
)

type PasswdCb func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error)
type PubKeyCb func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error)
type KeyIntCb func(conn ssh.ConnMetadata, cli ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error)
type AuthLogCb func(conn ssh.ConnMetadata, method string, err error)
type SvrLogFunc func(level int, msg string, err error)

type ZSshRpcServerCfg struct {
	HostKey                     ssh.Signer
	NoClientAuth                bool
	MaxAuthTries                int
	PasswordCallback            PasswdCb
	PublicKeyCallback           PubKeyCb
	KeyboardInteractiveCallback KeyIntCb
	AuthLogCallback             AuthLogCb
	ServerLogger                SvrLogFunc
	Support_AES128_CBC          bool
	OperationHandler            ZSshRpcOperationHandler
	IOBlockSize                 int
}

func NewDefaultZSshRpcServerCfg(host_key_store_path string) (*ZSshRpcServerCfg, error) {
	hostkey, err := GetOrGenerateHostKey(host_key_store_path, 2028)
	if err != nil {
		return nil, err
	}
	return &ZSshRpcServerCfg{
		HostKey:                     hostkey,
		NoClientAuth:                false,
		MaxAuthTries:                5,
		PasswordCallback:            defaultPasswordCallback,
		PublicKeyCallback:           defaultPubkeyCallback,
		KeyboardInteractiveCallback: defaultKeyIntCallback,
		AuthLogCallback:             defaultAuthLogCb,
		ServerLogger:                defaultServerLoggerFunc,
		Support_AES128_CBC:          true,
		OperationHandler:            &defaultOperationHandler{},
		IOBlockSize:                 1024,
	}, nil
}

func defaultPasswordCallback(conn ssh.ConnMetadata, passwd []byte) (*ssh.Permissions, error) {
	return nil, errors.New("method not support")
}

func defaultPubkeyCallback(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	return nil, errors.New("method not support")
}

func defaultKeyIntCallback(conn ssh.ConnMetadata, cli ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	return nil, errors.New("method not support")
}

func defaultAuthLogCb(conn ssh.ConnMetadata, method string, err error) {
}

type defaultOperationHandler struct {
}

func (this *defaultOperationHandler) HandleOperation(request ZSshRpcOperationRequest) ZSshRpcOperationResponse {
	return ZSshRpcOperationResponse{
		StatusCode:   ResponseStatusCode_NOT_FOUND,
		ResponseJSON: "{}",
	}
}

func defaultServerLoggerFunc(level int, msg string, err error) {
}

func (this *ZSshRpcServerCfg) UsePasswordAuth(callback PasswdCb) {
	this.PasswordCallback = callback
}

func (this *ZSshRpcServerCfg) UsePublicKeyAuth(callback PubKeyCb) {
	this.PublicKeyCallback = callback
}
func (this *ZSshRpcServerCfg) UseKeyboardInteractiveAuth(callback KeyIntCb) {
	this.KeyboardInteractiveCallback = callback
}
func (this *ZSshRpcServerCfg) UseAuthLogger(callback AuthLogCb) {
	this.AuthLogCallback = callback
}

func (this *ZSshRpcServerCfg) UseLogger(logfunc SvrLogFunc) {
	this.ServerLogger = logfunc
}

func (this *ZSshRpcServerCfg) SetHandler(handler ZSshRpcOperationHandler) {
	this.OperationHandler = handler
}

func GetOrGenerateHostKey(key_store_path string, key_length_to_gen int) (ssh.Signer, error) {
	if IsFileExists(key_store_path) {
		hostkeydata, err := ioutil.ReadFile(key_store_path)
		if err != nil {
			return generateHostKey(key_store_path, key_length_to_gen)
		}
		hostkey, err := ssh.ParsePrivateKey(hostkeydata)
		if err != nil {
			return generateHostKey(key_store_path, key_length_to_gen)
		}
		return hostkey, nil
	} else {
		return generateHostKey(key_store_path, key_length_to_gen)
	}
}

func generateHostKey(keypath string, len int) (ssh.Signer, error) {
	prikey, err := rsa.GenerateKey(rand.Reader, len)
	if err != nil {
		return nil, err
	}
	privateKeyFile, err := os.Create(keypath)
	defer privateKeyFile.Close()
	if err != nil {
		return nil, err
	}
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(prikey)}
	if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil {
		return nil, err
	}
	hostkeydata, err := ioutil.ReadFile(keypath)
	if err != nil {
		return nil, err
	}
	hostkey, err := ssh.ParsePrivateKey(hostkeydata)
	if err != nil {
		return nil, err
	}
	return hostkey, nil
}