package SSHServer

import (
	"git.swzry.com/NSMCServerLauncher/Utils"
	"fmt"
	"git.swzry.com/NSMCServerLauncher/Logger"
	"crypto/rsa"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"os"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"net"
	"github.com/swzry/go.TSmap"
	"golang.org/x/crypto/ssh/terminal"
	"git.swzry.com/NSMCServerLauncher/Terminal"
	"golang.org/x/crypto/bcrypt"
)

var mainThreadBlockChan chan byte
var ServerConf *ssh.ServerConfig
var _clientListRawMap map[interface{}]interface{}
var ClientList TSmap.TSmap

func StartSSHServer(mtbc chan byte)  {
	mainThreadBlockChan = mtbc
	if r,err := Utils.PathExists(SSHServerConf.host_key_file); ((!r) || (err != nil)) {
		fmt.Fprintln(&Logger.Log.SysInfo,"Generating SSH Host Key......")
		fmt.Fprintln(&Logger.Log.SysInfo,"(It may takes a long time. Please wait.)")
		fmt.Println("Generating SSH Host Key......")
		fmt.Println("(It may takes a long time. Please wait.)")
		GenerateRSAKey(SSHServerConf.host_key_file, SSHServerConf.defult_key_length)
	}
	ServerConf = &ssh.ServerConfig{
		NoClientAuth: false,
		MaxAuthTries: SSHServerConf.max_auth_tries,
		PasswordCallback: SSHPasswordCallback,
		PublicKeyCallback: SSHPublicKeyCallback,
		AuthLogCallback: SSHAuthLogCallback,
		ServerVersion: "SSH-2.0-NSMCServerLauncher-SSH",
	}
	keyBytes, err := ioutil.ReadFile(SSHServerConf.host_key_file)
	if(err != nil){
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Load Host Key File: ", err)
		mainThreadBlockChan <- 1
		return
	}
	key, err := ssh.ParsePrivateKey(keyBytes)
	if(err != nil){
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Load Host Key: ", err)
		mainThreadBlockChan <- 1
		return
	}
	ServerConf.AddHostKey(key)
	listener,err := net.Listen("tcp", SSHServerConf.bind_addr)
	if(err != nil){
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Bind Address '",SSHServerConf.bind_addr,"': ", err)
		mainThreadBlockChan <- 1
		return
	}
	fmt.Fprintln(&Logger.Log.SysInfo,"Listening At '",listener.Addr().String(),"'.")
	_clientListRawMap = make(map[interface{}]interface{})
	ClientList = &TSmap.NewTSmap{
		ConMap:_clientListRawMap,
	}
	for{
		tcpConn,err := listener.Accept()
		if(err == nil){
			fmt.Fprintf(&Logger.Log.SSH,"New Client '%v' Entered.\n",tcpConn.RemoteAddr())
			_, schan, reqchan, err := ssh.NewServerConn(tcpConn, ServerConf)
			if(err != nil){
				fmt.Fprintf(&Logger.Log.SSH,"Failed Handle Client '%v': %v\n",tcpConn.RemoteAddr(),err)
				tcpConn.Close()
				continue
			}else {
				ClientList.Set(tcpConn,&Utils.ClientConnection{
					Channels: &TSmap.NewTSmap{
						ConMap:make(map[interface{}]interface{}),
					},
				})
				go ssh.DiscardRequests(reqchan)
				go handleChannels(schan, tcpConn)
			}
		}
	}
}

func handleChannels(ch <-chan ssh.NewChannel, conn net.Conn)  {
	cnt := 0
	chlist,ok := ClientList.Get(conn)
	if(!ok){
		fmt.Fprintf(&Logger.Log.SSH,"Failed Handle Client '%v': Channels Not Found.\n",conn.RemoteAddr())
	}
	for newchan := range ch {
		cnt++;
		fmt.Fprintf(&Logger.Log.SSH,"Client '%v', Channels %v, Handling.\n",conn.RemoteAddr(),cnt)
		go handleChannel(newchan,conn,cnt,chlist.(*Utils.ClientConnection))
	}
}

func handleChannel(nch ssh.NewChannel,conn net.Conn,num int,chlist *Utils.ClientConnection)  {
	if(nch.ChannelType() == "session"){
		ch,req,err := nch.Accept()
		if(err != nil){
			fmt.Fprintf(&Logger.Log.SSH,"Client '%v', Channels %v, Failed Handling : %v\n",conn.RemoteAddr(),num,err)
			return
		}else{
			r:= <-req
			r.Reply(true,nil)
			chlist.Channels.Set(num,Utils.AvaliableChannel{
				Channel: ch,
				Term:    terminal.NewTerminal(ch,"NSMC >"),
				UserContext: Utils.UserContextType{
					PWD: "/",
				},
			})
			go Terminal.HandlerTerminal(ClientList,conn, num)
		}
	}else{
		nch.Reject(ssh.UnknownChannelType, "Unknown Channel Type")
		fmt.Fprintf(&Logger.Log.SSH,"Client '%v', Channels %v, Rejected : Unknown Channel Type '%v'\n",conn.RemoteAddr(),num,nch.ChannelType())
	}
}

func SSHPasswordCallback(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error)  {
	val, ok := SSHServerConf.passwd[conn.User()]
	if(!ok){
		fmt.Fprintf(&Logger.Log.SSH,"Client '%v' Auth Failed : Unknown User '%v'\n",conn.RemoteAddr(),conn.User())
		return nil,fmt.Errorf("Password Error or User '%v' Does Not Exist..",conn.User())
	}
	err := bcrypt.CompareHashAndPassword([]byte(val),password)
	if(err != nil){
		fmt.Fprintf(&Logger.Log.SSH,"Client '%v', User '%v', Auth Failed : Password Error, Detail: '%v'\n",conn.RemoteAddr(),conn.User(),err)
		return nil,fmt.Errorf("Password Error or User '%v' Does Not Exist..",conn.User())
	}
	return nil,nil
}

func SSHPublicKeyCallback(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error){
	return nil,fmt.Errorf("PublicKey Auth Not Support Yet.")
}

func SSHAuthLogCallback(conn ssh.ConnMetadata, method string, err error){
	fmt.Fprintf(&Logger.Log.SSH,"AuthFailed. Info: client='%v', user='%v', method='%v'. Err: %v\n",conn.RemoteAddr(),conn.User(),method,err)
}

func GenerateRSAKey(keyfn string,keyLength int) {
	privateKey, err := rsa.GenerateKey(rand.Reader, keyLength)
	if err != nil {
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Generate RSA Private Key: ", err)
		mainThreadBlockChan <- 1
		return
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	file, err := os.Create(keyfn)
	if err != nil {
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Create Key File: ", err)
		mainThreadBlockChan <- 1
		return
	}
	err = pem.Encode(file, block)
	if err != nil {
		fmt.Fprintln(&Logger.Log.SysFatal,"Failed to Encode Key: ", err)
		mainThreadBlockChan <- 1
		return
	}
	fmt.Fprintln(&Logger.Log.SysInfo,"SSH Host Key Generated Successfully..")
	fmt.Println("SSH Host Key Generated Successfully..")
}