package cmd

import (
	"Yasso/config"
	"fmt"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
	"io/ioutil"
	"os"
	"path"
)

/*
模块完成时间2021年12月28日，主要用于ssh爆破，连接，扫描22端口
*/

var SshCmd = &cobra.Command{
	Use: "ssh",
	Short: "SSH burst and SSH extend tools (support proxy)",
	Run: func(cmd *cobra.Command, args []string) {
		if Hosts == "" && ConnHost == "" {
			_ = cmd.Help()
		}else{
			BruteSshByUser()
		}
	},
}

func init(){
	SshCmd.Flags().StringVar(&ConnHost,"hostname","","Open an interactive SSH at that address(brute param need false)")
	SshCmd.Flags().StringVar(&LoginUser,"user","","Login ssh username")
	SshCmd.Flags().StringVar(&LoginPass,"pass","","Login ssh password")
	SshCmd.Flags().StringVar(&LoginPublicKey,"key","","ssh public key path")
}



func BruteSshByUser() {
	if BrutePort == 0 {
		BrutePort = 22
	}
	var ips []string
	var err error
	if Hosts != "" && ConnHost == ""{
		ips,err = ResolveIPS(Hosts)
		if err != nil {
			Println(fmt.Sprintf("resolve hosts address failed %v",err))
			return
		}
		if BruteFlag == true {
			users,pass := ReadTextToDic("ssh",UserDic,PassDic)
			Println(Clearln+"[*] Brute Module [ssh]")
			Println(fmt.Sprintf(Clearln + "[*] Have [user:%v] [pass:%v] [request:%v]",len(users),len(pass),len(users) * len(pass) * len(ips)))
			SwitchBurp("ssh",users,pass,ips,BrutePort,Runtime,TimeDuration,"")
		}else{
			Println(Clearln + "[*] May be you want to brute? try to add --crack")
		}
	}
	if ConnHost != "" && Hosts == "" && (LoginUser != "" && (LoginPass != "" || LoginPublicKey != "")) && BruteFlag != true {
		if LoginUser != ""  && LoginPass != ""  {
			client,status,err := SshConnByUser(config.HostIn{Host: ConnHost,Port: BrutePort,TimeOut: TimeDuration},LoginUser,LoginPass)
			if err != nil {
				Println(fmt.Sprintf(Clearln + "[-] Login ssh failed %v",err))
				return
			}
			if status == true {
				//认证成功
				SshLogin(client)
			}else {
				Println(Clearln + "[-] The username or password is incorrect")
				return
			}
		}
		if LoginPublicKey != "" && LoginUser != "" {
			client,status, err := sshConnByKey(config.HostIn{Host: ConnHost,Port: BrutePort,TimeOut: TimeDuration,PublicKey: LoginPublicKey},LoginUser)
			if err != nil {
				Println(fmt.Sprintf(Clearln + "[-] Login ssh failed %v",err))
				return
			}
			if status == true {
				//认证成功
				SshLogin(client)
				return
			}else {
				Println(Clearln + "[-] The username or password is incorrect")
				return
			}
		}
	}
	if Hosts == "" && ConnHost != "" && BruteFlag == false && (LoginUser == "" || LoginPublicKey == ""){
		Println(Clearln + "[*] May be you want login ssh? try to add user and (user' key) or (user' pass)")
		return
	}
}

func SshConnByUser(info config.HostIn,user,pass string)(*ssh.Client,bool,error){
	// 走socks5代理的ssh连接
	sshConfig := &ssh.ClientConfig{User: user,Auth: []ssh.AuthMethod{ssh.Password(pass)},HostKeyCallback: ssh.InsecureIgnoreHostKey(),Timeout:info.TimeOut}
	con, err := GetConn(fmt.Sprintf("%v:%v",info.Host,info.Port),info.TimeOut)
	if err != nil {
		return nil,false,err
	}
	c,ch,re,err := ssh.NewClientConn(con,fmt.Sprintf("%v:%v",info.Host,info.Port),sshConfig)
	if err != nil {
		return nil,false,err
	}
	return ssh.NewClient(c,ch,re),true,err
}

func sshConnByKey(info config.HostIn,user string)(*ssh.Client,bool,error){
	var (
		err error
		HomePath string
		key []byte
	)
	switch  {
	case info.PublicKey == "":
		HomePath, err = os.UserHomeDir()
		if err != nil {
			return nil,false,err
		}
		key, err = ioutil.ReadFile(path.Join(HomePath,".ssh","id_rsa"))
		if err != nil {
			return nil,false,err
		}
	case info.PublicKey != "":
		key, err = ioutil.ReadFile(info.PublicKey)
		if err != nil {
			return nil,false,err
		}
	}
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil,false,err
	}
	sshConfig := &ssh.ClientConfig{
		User: user,
		Auth:[]ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		Timeout: info.TimeOut,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}
	con, err := GetConn(fmt.Sprintf("%v:%v",info.Host,info.Port),info.TimeOut)
	if err != nil {
		return nil,false,err
	}

	c,ch,re,err := ssh.NewClientConn(con,fmt.Sprintf("%v:%v",info.Host,info.Port),sshConfig)
	if err != nil {
		return nil,false,err
	}
	return ssh.NewClient(c,ch,re),true,err
}


// ssh 完全交互式登陆

func SshLogin(client *ssh.Client) {
	defer client.Close()
	session, err := client.NewSession()
	if err != nil {
		Println(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 {
		Println(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 {
		Println(fmt.Sprintf("Session Request new xterm failed %v",err))
		return
	}
	if err = session.Shell(); err != nil {
		Println(fmt.Sprintf("Session start shell failed %v",err))
		return
	}
	if err = session.Wait();err != nil {
		Println(fmt.Sprintf("Session wait failed %v",err))
		return
	}
}