package main

import (
    "fmt"
    "net/http"
    "os"
    "strings"
)

var strValues string
var strUvalues string
var password string

func init() {
    password = os.Getenv("PASSWORD")
    if password == "" {
       password = "123456"
    }
}

func filterCookieValues(cookieValues string, keepKeys []string) string {
    newCookieValues := ""
    pairs := strings.Split(cookieValues, ";")
    if len(pairs) < 2 {
       return newCookieValues
    }
    for _, pair := range pairs {
       parts := strings.Split(pair, "=")
       key := strings.TrimSpace(parts[0])
       value := strings.Join(parts[1:], "=")
       if contains(keepKeys, key) {
          newCookieValues += key + "=" + value + "; "
       }
    }
    newCookieValues = newCookieValues[:len(newCookieValues)-2]
    return newCookieValues
}

func contains(keys []string, key string) bool {
    for _, k := range keys {
       if k == key {
          return true
       }
    }
    return false
}

func handleSET(w http.ResponseWriter, req *http.Request) {
    pwd := req.URL.Query().Get("pwd")
    if pwd == "" || pwd!= password {
       w.WriteHeader(http.StatusUnauthorized)
       w.Write([]byte("Invalid password"))
       return
    }
    keepKeys := []string{"_U", "MUID", "KievRPSSecAuth", "cct", "buid", "ak_bmsc", "bm_sv", "_RwBf", "SRCHHPGUSR", "WLS"}
    keepKeysU := []string{"_U", "WLS"}
    cookieValues := req.Header.Get("Cookie-Values")
    setValue := filterCookieValues(cookieValues, keepKeys)
    getUValue := filterCookieValues(cookieValues, keepKeysU)

    if setValue!= "" {
       strValues = setValue
       if getUValue!= "" &&!strings.Contains(strUvalues, getUValue) {
          strUvalues += ";" + getUValue
       }
       w.Write([]byte("Set value successfully"))
    } else {
       w.WriteHeader(http.StatusBadRequest)
       w.Write([]byte("No Cookie-Values in header"))
    }
}

func handleGET(w http.ResponseWriter, req *http.Request) {
    pwd := req.URL.Query().Get("pwd")
    if pwd == "" || pwd!= password {
       w.WriteHeader(http.StatusUnauthorized)
       w.Write([]byte("Invalid password"))
       return
    }
    result := `{"result": {"cookies": "` + strValues + `"}}`
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(fmt.Sprintf("%v", result)))
}

func handleCLS(w http.ResponseWriter, req *http.Request) {
    pwd := req.URL.Query().Get("pwd")
    if pwd == "" || pwd!= password {
       w.WriteHeader(http.StatusUnauthorized)
       w.Write([]byte("Invalid password"))
       return
    }
    replacedStr := strings.ReplaceAll(strUvalues, ";", "<br>")

    strValues = ""
    strUvalues = ""

    w.Write([]byte("Clear value successfully\n" + replacedStr))
}

func handleHisU(w http.ResponseWriter, req *http.Request) {
    pwd := req.URL.Query().Get("pwd")
    if pwd == "" || pwd!= password {
       w.WriteHeader(http.StatusUnauthorized)
       w.Write([]byte("Invalid password"))
       return
    }
    replacedStr := strings.ReplaceAll(strUvalues, ";", "<br>")
    w.Write([]byte("Ukey History:\n" + replacedStr))
}

func handleRoot(w http.ResponseWriter, req *http.Request) {
    w.Write([]byte("Please visit /SET, /GET, or /CLS with?pwd=xxxxxx"))
}

// 新增处理 /q 的函数
func handleQ(w http.ResponseWriter, req *http.Request) {
    url := req.URL.Query().Get("url")
    if url == "" {
       w.WriteHeader(http.StatusBadRequest)
       w.Write([]byte("URL is required"))
       return
    }
    keepKeys := []string{"_U", "MUID", "KievRPSSecAuth", "cct", "buid", "ak_bmsc", "bm_sv", "_RwBf", "SRCHHPGUSR", "WLS"}
    client := &http.Client{}
    req, err := http.NewRequest("GET", url, nil)
    if err!= nil {
       w.WriteHeader(http.StatusInternalServerError)
       w.Write([]byte(err.Error()))
       return
    }
    req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0")
    req.Header.Set("Accept-Language", "zh-CN,zh;q=0.9")
    resp, err := client.Do(req)
    if err!= nil {
       w.WriteHeader(http.StatusInternalServerError)
       w.Write([]byte(err.Error()))
       return
    }
    defer resp.Body.Close()
    setCookieArray := resp.Header["Set-Cookie"]
    cookies := ""
    if setCookieArray!= nil {
       // 只保留键值对，并替换 & 为 _
       cookies = setCookieArray[0]
       pairs := strings.Split(cookies, ";")
       keyValuePair := pairs[0]
       keyValuePair = strings.ReplaceAll(keyValuePair, "&", "_")
       cookies = keyValuePair
    }
    setValue := filterCookieValues(cookies, keepKeys)
    result := `{"result": {"cookies": "` + setValue + `"}}`
    w.Header().Set("Content-Type", "application/json")
    w.Write([]byte(fmt.Sprintf("%v", result)))
}

func main() {
    port := ":7860"
    http.HandleFunc("/SET", handleSET)
    http.HandleFunc("/GET", handleGET)
    http.HandleFunc("/CLS", handleCLS)
    http.HandleFunc("/HisU", handleHisU)
    // 新增处理 /q 的路由
    http.HandleFunc("/q", handleQ)
    http.HandleFunc("/", handleRoot)


    fmt.Printf("Server is running on port %s\n", port)
    http.ListenAndServe(port, nil)
}