package ssh

import (
	config2 "Yasso/config"
	"Yasso/core/logger"
	"Yasso/core/plugin"
	"Yasso/pkg/exploit/config"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
	"os"
	"time"
)

func ExploitSSH(exploits config.Exploits, key string) {
	var SshConn *ssh.Client
	var status bool
	var err error
	if key == "" {
		SshConn, status, err = plugin.SshConnByUser(config2.ServiceConn{
			Hostname: exploits.Hostname,
			Port:     exploits.Port,
			Timeout:  1000 * time.Millisecond,
		}, exploits.User, exploits.Pass)
	} else {
		SshConn, status, err = plugin.SshConnByKey(config2.ServiceConn{
			Hostname:  exploits.Hostname,
			Port:      exploits.Port,
			Timeout:   1000 * time.Millisecond,
			PublicKey: key,
		}, exploits.User)
	}
	if err != nil || status == false {
		logger.Fatal("exploit ssh has an error conn to ssh failed")
		return
	}
	loginSSH(SshConn) // 连接到ssh
}

func loginSSH(client *ssh.Client) {
	defer client.Close()
	session, err := client.NewSession()
	if err != nil {
		logger.Fatal(fmt.Sprintf("new ssh session failed %v", err))
		return
	}
	defer session.Close()
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
		ssh.VSTATUS:       1,
	}
	fd := int(os.Stdin.Fd())
	oldState, err := terminal.MakeRaw(fd)
	if err != nil {
		logger.Fatal(fmt.Sprintf("terminal failed %v", err))
	}
	defer terminal.Restore(fd, oldState)
	w, h, err := terminal.GetSize(fd)
	if err = session.RequestPty("xterm-256color", h, w, modes); err != nil {
		logger.Fatal(fmt.Sprintf("Session Request new xterm failed %v", err))
		return
	}
	if err = session.Shell(); err != nil {
		logger.Fatal(fmt.Sprintf("Session start shell failed %v", err))
		return
	}
	if err = session.Wait(); err != nil {
		logger.Fatal(fmt.Sprintf("Session wait failed %v", err))
		return
	}
}