mirror of
https://gitee.com/chuangxxt/share-copilot
synced 2025-04-16 10:13:26 +00:00
317 lines
8.2 KiB
Go
317 lines
8.2 KiB
Go
package main
|
||
|
||
import (
|
||
"crypto/tls"
|
||
"encoding/json"
|
||
"fmt"
|
||
"github.com/fatih/color"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/go-resty/resty/v2"
|
||
"io"
|
||
"log"
|
||
"math/rand"
|
||
"net/http"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
func main() {
|
||
//初始化配置文件
|
||
configFile = initConfig()
|
||
// 创建Gin引擎
|
||
gin.SetMode(gin.ReleaseMode)
|
||
gin.DefaultWriter = io.Discard
|
||
engine := gin.New()
|
||
// 设置信任的代理
|
||
err := engine.SetTrustedProxies([]string{"127.0.0.1"})
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
// 初始化有效的token列表
|
||
initValidTokenList()
|
||
// 定义路由
|
||
domainDefault := engine.Group("/", DomainMiddleware(configFile.Server.Domain))
|
||
domainDefault.GET("/copilot_internal/v2/token", getGithubToken())
|
||
// 初始化服务器
|
||
initServer(engine)
|
||
// 显示信息
|
||
showMsg()
|
||
}
|
||
|
||
// 初始化服务器
|
||
func initServer(engine *gin.Engine) {
|
||
// 配置支持的应用程序协议
|
||
server := &http.Server{
|
||
Addr: ":443",
|
||
TLSConfig: &tls.Config{
|
||
NextProtos: []string{"http/1.1", "http/1.2", "http/2"}, // 支持的应用程序协议列表
|
||
},
|
||
Handler: engine,
|
||
}
|
||
// 启动Gin服务器并监听端口
|
||
listenAddress := configFile.Server.Host + ":" + strconv.Itoa(configFile.Server.Port)
|
||
go func() {
|
||
if configFile.Server.Port != 443 {
|
||
err := engine.Run(listenAddress)
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
} else {
|
||
err := server.ListenAndServeTLS(configFile.Server.CertPath, configFile.Server.KeyPath)
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// 初始化有效的token列表
|
||
func initValidTokenList() {
|
||
//为了安全起见，应该等待请求完成并处理其响应。
|
||
var wg sync.WaitGroup
|
||
for _, token := range configFile.CopilotConfig.Token {
|
||
wg.Add(1)
|
||
go func(token string) {
|
||
defer wg.Done()
|
||
if getGithubApi(token) {
|
||
validTokenList[token] = true
|
||
}
|
||
}(token)
|
||
}
|
||
wg.Wait()
|
||
}
|
||
|
||
// 获取Github的token
|
||
func getGithubToken() gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 请求计数
|
||
requestCount++
|
||
// 如果配置了verification，则需要获取请求头中的Authorization令牌
|
||
if configFile.Verification != "" {
|
||
token := c.GetHeader("Authorization")
|
||
tokenStr := strings.ReplaceAll(token, " ", "")
|
||
configCert := strings.ReplaceAll(configFile.Verification, " ", "")
|
||
if tokenStr != "token"+configCert {
|
||
// 拒绝不符合Verification的请求
|
||
badRequest(c)
|
||
return
|
||
}
|
||
}
|
||
//从有效的token列表中随机获取一个token
|
||
token := getRandomToken(validTokenList)
|
||
//判断tokenMap 里的token是否存在
|
||
if _, exists := tokenMap[token]; exists {
|
||
respDataMap := tokenMap[token]
|
||
//判断时间戳key是否存在
|
||
if _, exists := respDataMap["expires_at"]; exists {
|
||
// 获取当前时间的Unix时间戳
|
||
currentTime := time.Now().Unix()
|
||
if expiresAt, ok := respDataMap["expires_at"].(float64); ok {
|
||
// 判断expires_at是否已经过期
|
||
expiresAtInt64 := int64(expiresAt)
|
||
//提前一分钟请求
|
||
if expiresAtInt64 > currentTime+60 {
|
||
//fmt.Println("\n未过期无需请求")
|
||
proxyResp(c, tokenMap[token])
|
||
} else {
|
||
//fmt.Println("\n已过期重新请求")
|
||
if getGithubApi(token) {
|
||
proxyResp(c, tokenMap[token])
|
||
return
|
||
} else {
|
||
badRequest(c)
|
||
}
|
||
}
|
||
} else {
|
||
badRequest(c)
|
||
}
|
||
} else {
|
||
//tokenMap里的token对应的返回体不存在expires_at
|
||
if getGithubApi(token) {
|
||
proxyResp(c, tokenMap[token])
|
||
return
|
||
} else {
|
||
badRequest(c)
|
||
}
|
||
}
|
||
} else {
|
||
//不存在则githubApi发送请求，并存到tokenMap
|
||
if getGithubApi(token) {
|
||
proxyResp(c, tokenMap[token])
|
||
return
|
||
} else {
|
||
badRequest(c)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 请求错误
|
||
func badRequest(c *gin.Context) {
|
||
c.JSON(http.StatusBadRequest, gin.H{
|
||
"message": "Bad credentials",
|
||
"documentation_url": "https://docs.github.com/rest"})
|
||
}
|
||
|
||
// 本服务器响应
|
||
func proxyResp(c *gin.Context, respDataMap map[string]interface{}) {
|
||
// 将map转换为JSON字符串
|
||
responseJSON, err := json.Marshal(respDataMap)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{"error": "JSON marshaling error"})
|
||
}
|
||
// 请求成功统计
|
||
successCount++
|
||
// 将JSON字符串作为响应体返回
|
||
c.Header("Content-Type", "application/json")
|
||
c.String(http.StatusOK, string(responseJSON))
|
||
}
|
||
|
||
// 请求githubApi
|
||
func getGithubApi(token string) bool {
|
||
githubApiCount++
|
||
// 设置请求头
|
||
headers := map[string]string{
|
||
"Authorization": "token " + token,
|
||
/*"editor-version": c.GetHeader("editor-version"),
|
||
"editor-plugin-version": c.GetHeader("editor-plugin-version"),
|
||
"user-agent": c.GetHeader("user-agent"),
|
||
"accept": c.GetHeader("accept"),
|
||
"accept-encoding": c.GetHeader("accept-encoding"),*/
|
||
}
|
||
// 发起GET请求
|
||
response, err := resty.New().R().
|
||
SetHeaders(headers).
|
||
Get(configFile.CopilotConfig.GithubApiUrl)
|
||
if err != nil {
|
||
// 处理请求错误
|
||
return false
|
||
}
|
||
// 判断响应状态码
|
||
if response.StatusCode() == http.StatusOK {
|
||
// 响应状态码为200 OK
|
||
respDataMap := map[string]interface{}{}
|
||
err = json.Unmarshal(response.Body(), &respDataMap)
|
||
if err != nil {
|
||
// 处理JSON解析错误
|
||
return false
|
||
}
|
||
//token map
|
||
tokenMap[token] = respDataMap
|
||
return true
|
||
} else {
|
||
// 处理其他状态码
|
||
delete(validTokenList, token)
|
||
return false
|
||
}
|
||
}
|
||
|
||
// 初始化配置文件
|
||
func initConfig() Config {
|
||
// 读取配置文件
|
||
exePath, err := os.Executable()
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
// 获取执行文件所在目录
|
||
exeDir := filepath.Dir(exePath)
|
||
configFile, err := os.Open(exeDir + "/config.json")
|
||
if err != nil {
|
||
panic("file \"./config.json\" not found")
|
||
}
|
||
defer func(configFile *os.File) {
|
||
err := configFile.Close()
|
||
if err != nil {
|
||
panic("close file \"./config.json\" err")
|
||
}
|
||
}(configFile)
|
||
decoder := json.NewDecoder(configFile)
|
||
config := Config{}
|
||
err = decoder.Decode(&config)
|
||
if err != nil {
|
||
panic("config format err")
|
||
}
|
||
return config
|
||
}
|
||
|
||
// 重置请求计数
|
||
func resetRequestCount() {
|
||
requestCountMutex.Lock()
|
||
defer requestCountMutex.Unlock()
|
||
requestCount = 0
|
||
successCount = 0
|
||
}
|
||
|
||
// 从map中随机获取一个key
|
||
func getRandomToken(m map[string]bool) string {
|
||
keys := make([]string, 0, len(m))
|
||
for k := range m {
|
||
keys = append(keys, k)
|
||
}
|
||
if len(keys) == 0 {
|
||
return "" // 返回空字符串或处理其他错误情况
|
||
}
|
||
r := rand.New(rand.NewSource(time.Now().UnixNano()))
|
||
randomIndex := r.Intn(len(keys))
|
||
return keys[randomIndex]
|
||
}
|
||
|
||
// DomainMiddleware 域名中间件
|
||
func DomainMiddleware(domain string) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 检查域名是否匹配
|
||
requestDomain := strings.Split(c.Request.Host, ":")[0]
|
||
if requestDomain == domain || requestDomain == "127.0.0.1" {
|
||
c.Next()
|
||
} else {
|
||
c.String(403, "Forbidden")
|
||
c.Abort()
|
||
}
|
||
}
|
||
}
|
||
|
||
// 显示信息
|
||
func showMsg() {
|
||
var url = ""
|
||
if configFile.Server.Port == 80 {
|
||
url = "http://" + configFile.Server.Domain
|
||
} else if configFile.Server.Port == 443 {
|
||
url = "https://" + configFile.Server.Domain
|
||
} else {
|
||
url = "http://" + configFile.Server.Domain + ":" + strconv.Itoa(configFile.Server.Port)
|
||
}
|
||
var jetStr = color.WhiteString("[Jetbrains]")
|
||
var vsStr = color.WhiteString("[Vscode]")
|
||
var valid = color.WhiteString("[Valid tokens]")
|
||
|
||
fmt.Println(jetStr + ": " + color.HiBlueString(url+"/copilot_internal/v2/token"))
|
||
fmt.Println(vsStr + ": " + color.HiBlueString(url))
|
||
fmt.Println(valid + ": " + color.HiBlueString(strconv.Itoa(len(validTokenList))))
|
||
fmt.Println(color.WhiteString("-----------------------------------------------------------------------"))
|
||
for {
|
||
requestCountMutex.Lock()
|
||
sCount := successCount
|
||
tCount := requestCount
|
||
gCount := githubApiCount
|
||
requestCountMutex.Unlock()
|
||
currentTime := time.Now().Format("2006-01-02 15:04:05")
|
||
if "00:00:00" == currentTime {
|
||
resetRequestCount()
|
||
}
|
||
var s2 = color.WhiteString("[Succeed]")
|
||
var s3 = color.WhiteString("[Failed]")
|
||
var s4 = color.WhiteString("[GithubApi]")
|
||
// 打印文本
|
||
fmt.Printf("\033[G%s - %s: %s %s: %s %s: %s ",
|
||
color.HiYellowString(currentTime),
|
||
s2, color.GreenString(strconv.Itoa(sCount)),
|
||
s3, color.RedString(strconv.Itoa(tCount-sCount)),
|
||
s4, color.CyanString(strconv.Itoa(gCount)))
|
||
time.Sleep(1 * time.Second) //
|
||
}
|
||
}
|