package exploit

import (
	"Yasso/core/logger"
	"Yasso/pkg/exploit/config"
	"Yasso/pkg/exploit/ldap"
	"Yasso/pkg/exploit/mssql"
	"Yasso/pkg/exploit/redis"
	"Yasso/pkg/exploit/ssh"
	"Yasso/pkg/exploit/sunlogin"
	"Yasso/pkg/exploit/winrm"
	"github.com/spf13/cobra"
)

type ExpFlags struct {
	Hostname string
	Port     int
	User     string
	Pass     string
	KeyFile  string
	Rebound  string
	Command  string
	Method   int    // 每一个Exp的子方法
	Listen   string // 本地监听地址
	SoPath   string // so文件路径
	Filter   string // ldap的过滤器
	LdapCmd  bool   // ldap的查询命令显示
	LdapAll  bool   // 是否自动查询ldap
	LdapName string // Ldap的name属性
}

var mssqlFlag ExpFlags
var MssqlCmd = &cobra.Command{
	Use:   "mssql",
	Short: "Quick attacks on MSSQL services",
	Run: func(cmd *cobra.Command, args []string) {
		if mssqlFlag.Hostname == "" {
			_ = cmd.Help()
			return
		}
		switch mssqlFlag.Method {
		case 1:
			mssql.ExploitMssql(config.Exploits{
				Hostname: mssqlFlag.Hostname,
				Port:     mssqlFlag.Port,
				User:     mssqlFlag.User,
				Pass:     mssqlFlag.Pass,
			}, 1, mssqlFlag.Command)
		case 2:
			mssql.ExploitMssql(config.Exploits{
				Hostname: mssqlFlag.Hostname,
				Port:     mssqlFlag.Port,
				User:     mssqlFlag.User,
				Pass:     mssqlFlag.Pass,
			}, 2, mssqlFlag.Command)
		case 3:
			mssql.ExploitMssql(config.Exploits{
				Hostname: mssqlFlag.Hostname,
				Port:     mssqlFlag.Port,
				User:     mssqlFlag.User,
				Pass:     mssqlFlag.Pass,
			}, 3, mssqlFlag.Command)
		case 4:
			mssql.ExploitMssql(config.Exploits{
				Hostname: mssqlFlag.Hostname,
				Port:     mssqlFlag.Port,
				User:     mssqlFlag.User,
				Pass:     mssqlFlag.Pass,
			}, 4, mssqlFlag.Command)
		default:
			logger.Fatal("not found exploit method")
			return
		}
	},
}

var sshFlag ExpFlags
var SshCmd = &cobra.Command{
	Use:   "ssh",
	Short: "Quick attacks on SSH services",
	Run: func(cmd *cobra.Command, args []string) {
		if sshFlag.Hostname == "" {
			_ = cmd.Help()
			return
		}
		ssh.ExploitSSH(config.Exploits{
			Hostname: sshFlag.Hostname,
			Port:     sshFlag.Port,
			User:     sshFlag.User,
			Pass:     sshFlag.Pass,
		}, sshFlag.KeyFile)
	},
}

var winrmFlag ExpFlags
var WinRmCmd = &cobra.Command{
	Use:   "winrm",
	Short: "Quick attacks on WinRM services",
	Run: func(cmd *cobra.Command, args []string) {
		if winrmFlag.Hostname == "" {
			_ = cmd.Help()
			return
		}
		winrm.ExploitWinRM(config.Exploits{
			Hostname: winrmFlag.Hostname,
			Port:     winrmFlag.Port,
			User:     winrmFlag.User,
			Pass:     winrmFlag.Pass,
		}, winrmFlag.Command, winrmFlag.Method)
	},
}

var redisFlag ExpFlags
var RedisCmd = &cobra.Command{
	Use:   "redis",
	Short: "Quick attacks on Redis services",
	Run: func(cmd *cobra.Command, args []string) {
		if redisFlag.Hostname == "" {
			_ = cmd.Help()
			return
		}
		redis.ExploitRedis(config.Exploits{
			Hostname: redisFlag.Hostname,
			Port:     redisFlag.Port,
			User:     "",
			Pass:     redisFlag.Pass,
		}, redisFlag.Method, redisFlag.Rebound, redisFlag.KeyFile, redisFlag.Listen, "")
	},
}

var sunLoginFlag ExpFlags
var SunLoginCmd = &cobra.Command{
	Use:   "sunlogin",
	Short: "Quick attacks on SunLogin services (RCE)",
	Run: func(cmd *cobra.Command, args []string) {
		if sunLoginFlag.Hostname == "" {
			_ = cmd.Help()
			return
		}
		if sunLoginFlag.Port == 0 {
			logger.Fatal("input sunlogin port")
			return
		} else {
			sunlogin.ExploitSunLogin(config.Exploits{
				Hostname: sunLoginFlag.Hostname,
				Port:     sunLoginFlag.Port,
				User:     "",
				Pass:     "",
			}, sunLoginFlag.Command)
		}
	},
}

var LdapReaperFlag ExpFlags

var LdapReaperCmd = &cobra.Command{
	Use:   "ldap",
	Short: "ldap single query with filter and fast automatic query",
	Run: func(cmd *cobra.Command, args []string) {
		if LdapReaperFlag.Hostname == "" || LdapReaperFlag.User == "" {
			_ = cmd.Help()
			return
		}
		if LdapReaperFlag.LdapCmd == true {
			ldap.ListLdapCommand()
			return
		} else {
			if LdapReaperFlag.Command != "" {
				LdapReaperFlag.LdapAll = false
			}
			ldap.LdapAuthAndQuery(LdapReaperFlag.Hostname, LdapReaperFlag.User, LdapReaperFlag.Pass, LdapReaperFlag.Command, LdapReaperFlag.Filter, LdapReaperFlag.LdapName, LdapReaperFlag.LdapAll)
		}
	},
}

func init() {
	MssqlCmd.Flags().StringVar(&mssqlFlag.Hostname, "host", "", "设置mssql连接主机地址")
	MssqlCmd.Flags().StringVar(&mssqlFlag.Command, "cmd", "", "执行的system命令")
	MssqlCmd.Flags().IntVar(&mssqlFlag.Port, "port", 1433, "设置mssql连接主机端口")
	MssqlCmd.Flags().StringVar(&mssqlFlag.User, "user", "sa", "设置连接的用户名")
	MssqlCmd.Flags().StringVar(&mssqlFlag.Pass, "pass", "", "设置连接的密码")
	MssqlCmd.Flags().IntVar(&mssqlFlag.Method, "method", 1, "设置exploit方法(.eg)\n[1][xp_cmdshell]\n[2][sp_oacreate]\n[3][install SharpSQLKit]\n[4][uninstall SharpSQLKit]")
	SshCmd.Flags().StringVar(&sshFlag.Hostname, "host", "", "设置ssh连接主机地址")
	SshCmd.Flags().StringVar(&sshFlag.KeyFile, "key", "", "设置ssh的连接密钥")
	SshCmd.Flags().StringVar(&sshFlag.User, "user", "root", "设置连接的用户名")
	SshCmd.Flags().StringVar(&sshFlag.Pass, "pass", "", "设置连接的密码")
	SshCmd.Flags().IntVar(&sshFlag.Port, "port", 22, "设置ssh的连接端口")
	RedisCmd.Flags().StringVar(&redisFlag.Hostname, "host", "", "设置redis主机的连接地址")
	RedisCmd.Flags().StringVar(&redisFlag.Rebound, "rebound", "", "设置redis定时计划反弹shell地址")
	RedisCmd.Flags().StringVar(&redisFlag.KeyFile, "key", "", "设置redis写入公钥的本地文件路径")
	RedisCmd.Flags().StringVar(&redisFlag.Listen, "listen", "127.0.0.1:8888", "设置redis主从服务本地监听")
	RedisCmd.Flags().StringVar(&redisFlag.Pass, "pass", "", "设置redis的连接密码")
	RedisCmd.Flags().StringVar(&redisFlag.SoPath, "so", "", "设置其他so文件路径")
	RedisCmd.Flags().IntVar(&redisFlag.Port, "port", 6379, "设置redis的连接端口")
	RedisCmd.Flags().IntVar(&redisFlag.Method, "method", 1, "设置exploit方法(.eg)\n[1][redis定时计划反弹shell]\n[2][redis公钥写入]\n[3][redis主从复制RCE(需要对方主机可以访问服务启动主机)]")
	WinRmCmd.Flags().StringVar(&winrmFlag.Hostname, "host", "", "设置winrm连接主机")
	WinRmCmd.Flags().StringVar(&winrmFlag.User, "user", "administrator", "设置winrm连接用户")
	WinRmCmd.Flags().StringVar(&winrmFlag.Pass, "pass", "", "设置winrm连接密码")
	WinRmCmd.Flags().StringVar(&winrmFlag.Command, "cmd", "whoami", "设置winrm执行的命令")
	WinRmCmd.Flags().IntVar(&winrmFlag.Port, "port", 5985, "设置winrm连接端口")
	WinRmCmd.Flags().IntVar(&winrmFlag.Method, "method", 1, "设置exploit方法(.eg)\n[1][winrm单命令执行,需配合cmd参数]\n[2][winrm正向shell方式执行]")
	SunLoginCmd.Flags().StringVar(&sunLoginFlag.Hostname, "host", "", "设置向日葵主机地址")
	SunLoginCmd.Flags().IntVar(&sunLoginFlag.Port, "port", 0, "设置向日葵端口")
	SunLoginCmd.Flags().StringVar(&sunLoginFlag.Command, "cmd", "whoami", "设置system命令")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.Hostname, "dc", "", "设置dc的主机名(FQDN)")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.User, "user", "", "设置域用户名称(.eg)[KLION\\Oadmin]")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.Pass, "pass", "", "设置域用户密码")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.Filter, "filter", "full-data", "设置过滤器,一般为full-data")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.Command, "cmd", "", "设置查询命令,可通过ldap-cmd查看")
	LdapReaperCmd.Flags().BoolVar(&LdapReaperFlag.LdapCmd, "ldap-cmd", false, "列出ldap可用的查询命令")
	LdapReaperCmd.Flags().BoolVar(&LdapReaperFlag.LdapAll, "ldap-all", true, "是否采用自动ldap查询(将查询默认ldap信息)")
	LdapReaperCmd.Flags().StringVar(&LdapReaperFlag.LdapName, "name", "", "域（成员,组,计算机）名称")
}