Merge branch 'refs/heads/feat-cmd'

This commit is contained in:
junleea 2025-03-05 19:59:58 +08:00
commit 6490779a2e
4 changed files with 32 additions and 14 deletions

View File

@ -35,14 +35,14 @@ func FindShellByID(id, uid uint) []Shell {
var result *gorm.DB var result *gorm.DB
if proto.Config.SERVER_SQL_LOG { if proto.Config.SERVER_SQL_LOG {
result = DB.Debug().Where("id = ? & auth_id = ?", id, uid).First(&shell) result = DB.Debug().Where("id = ? and auth_id = ?", id, uid).First(&shell)
} else { } else {
result = DB.Where("id = ? & auth_id = ?", id, uid).First(&shell) result = DB.Where("id = ? and auth_id = ?", id, uid).First(&shell)
} }
var res []Shell var res []Shell
res = append(res, shell) res = append(res, shell)
if result.Error != nil { if result.Error != nil {
log.Fatalf("FindShellByID failed: %v", result.Error) log.Printf("FindShellByID failed: %v", result.Error)
} }
return res return res
} }
@ -77,9 +77,24 @@ func UpdateShellByID(id, authId uint, shellName, shellContent string, status int
func FindShellWillRunByServer(server string, uid uint) []Shell { func FindShellWillRunByServer(server string, uid uint) []Shell {
var shells []Shell var shells []Shell
if proto.Config.SERVER_SQL_LOG { if proto.Config.SERVER_SQL_LOG {
DB.Debug().Exec("select * from shells where server = ? and auth_id = ? and status = 0 order by created_at desc limit 100", server, uid).Scan(&shells) result := DB.Debug().Where("server = ? AND auth_id = ? AND status = 0", server, uid).
Order("created_at DESC").
Limit(100).
Find(&shells)
if result.Error != nil {
// 若查询过程中出现错误,输出错误信息
log.Printf("Failed to query shells: %v\n", result.Error)
}
} else { } else {
DB.Exec("select * from shells where server = ? and auth_id = ? and status = 0 order by created_at desc limit 100", server, uid).Scan(&shells) //DB.Exec("select * from shells where server = ? and auth_id = ? and status = 0 order by created_at desc limit 100", server, uid).Scan(&shells)
result := DB.Where("server = ? AND auth_id = ? AND status = 0", server, uid).
Order("created_at DESC").
Limit(100).
Find(&shells)
if result.Error != nil {
// 若查询过程中出现错误,输出错误信息
log.Printf("Failed to query shells: %v\n", result.Error)
}
} }
return shells return shells
} }

12
main.go
View File

@ -105,7 +105,6 @@ func JWTAuthMiddleware() gin.HandlerFunc {
if tokenString == "" { if tokenString == "" {
tokenString = c.Query("token") tokenString = c.Query("token")
} }
//如果请求为login或register则不需要验证token //如果请求为login或register则不需要验证token
for k, _ := range proto.Url_map { for k, _ := range proto.Url_map {
if strings.Contains(c.Request.URL.Path, k) { if strings.Contains(c.Request.URL.Path, k) {
@ -134,9 +133,13 @@ func JWTAuthMiddleware() gin.HandlerFunc {
return return
} }
} }
//查看token是否在超级token中 //查看token是否在超级token中
if worker.IsContainSet("super_permission_tokens", tokenString) { if worker.IsContainSet("super_permission_tokens", tokenString) {
sId := c.Request.Header.Get("super_id") sId := c.Request.Header.Get("super_id")
if sId == "" {
sId = c.Query("super_id")
}
if sId == "" { if sId == "" {
c.AbortWithStatus(200) c.AbortWithStatus(200)
c.JSON(200, gin.H{ c.JSON(200, gin.H{
@ -146,12 +149,7 @@ func JWTAuthMiddleware() gin.HandlerFunc {
}) })
return return
} }
if sId == "" { log.Printf("req super_id:%s", sId)
sId = c.Query("super_id")
}
if sId == "" {
sId = "1"
}
id, _ := strconv.Atoi(sId) id, _ := strconv.Atoi(sId)
idFloat64 := float64(id) idFloat64 := float64(id)
//查看s_id类型 //查看s_id类型

View File

@ -57,10 +57,13 @@ func ShellWillRunFromServer() {
resp = append(resp, proto.UpdateShellReq{ID: v.ID, Server: v.Server, Status: 2, ShellResult: res}) resp = append(resp, proto.UpdateShellReq{ID: v.ID, Server: v.Server, Status: 2, ShellResult: res})
} }
} }
if len(resp) == 0 {
return
}
//返回执行结果 //返回执行结果
url := "https://" + proto.Config.MASTER_SERVER_DOMAIN + "/shell/update" url := "https://" + proto.Config.MASTER_SERVER_DOMAIN + "/shell/update"
var req proto.SyncUserShellResp var req proto.SyncUserShellResp
req.Token = worker.GetRedisSetMembers("super_permission_tokens")[0] req.Token = worker.GetRedisSetMembers("super_permission_tokens?super_id=1")[0]
req.Shells = resp req.Shells = resp
resp_data, err := worker.SyncDataFromMasterShellReq3(url, req) resp_data, err := worker.SyncDataFromMasterShellReq3(url, req)
if err != nil { if err != nil {
@ -78,7 +81,7 @@ func ShellWillRunFromServer() {
func GetShellWillRunFromMaster(server string) ([]dao.Shell, error) { func GetShellWillRunFromMaster(server string) ([]dao.Shell, error) {
master := proto.Config.MASTER_SERVER_DOMAIN master := proto.Config.MASTER_SERVER_DOMAIN
//发起请求获取待执行的shell //发起请求获取待执行的shell
url := "https://" + master + "/shell/server_will_run_list" url := "https://" + master + "/shell/server_will_run_list?super_id=1"
var req proto.SyncUserShellReq var req proto.SyncUserShellReq
req.Server = server req.Server = server
req.Token = worker.GetRedisSetMembers("super_permission_tokens")[0] req.Token = worker.GetRedisSetMembers("super_permission_tokens")[0]

View File

@ -211,6 +211,7 @@ func SyncDataFromMasterShellReq2(url string, data proto.SyncUserShellReq) ([]dao
return res, err return res, err
} }
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("token", data.Token)
//传输数据 //传输数据
if client == nil { if client == nil {
client = &http.Client{} client = &http.Client{}
@ -254,6 +255,7 @@ func SyncDataFromMasterShellReq3(url string, data proto.SyncUserShellResp) ([]pr
return res, err return res, err
} }
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("token", data.Token)
//传输数据 //传输数据
if client == nil { if client == nil {
client = &http.Client{} client = &http.Client{}