mirror of
https://github.com/XZB-1248/Spark
synced 2025-09-26 20:21:11 +08:00
399 lines
11 KiB
Go
399 lines
11 KiB
Go
package utility
|
|
|
|
import (
|
|
"Spark/modules"
|
|
"Spark/server/common"
|
|
"Spark/server/config"
|
|
"Spark/utils"
|
|
"Spark/utils/melody"
|
|
"bytes"
|
|
"context"
|
|
"crypto/aes"
|
|
"crypto/cipher"
|
|
"fmt"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/kataras/golog"
|
|
"net/http"
|
|
"os"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
type Sender func(pack modules.Packet, session *melody.Session) bool
|
|
|
|
// CheckForm checks if the form contains the required fields.
|
|
// Every request must contain connection UUID or device ID.
|
|
func CheckForm(ctx *gin.Context, form any) (string, bool) {
|
|
var base struct {
|
|
Conn string `json:"uuid" yaml:"uuid" form:"uuid"`
|
|
Device string `json:"device" yaml:"device" form:"device"`
|
|
}
|
|
if form != nil && ctx.ShouldBind(form) != nil {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return ``, false
|
|
}
|
|
if ctx.ShouldBind(&base) != nil || (len(base.Conn) == 0 && len(base.Device) == 0) {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return ``, false
|
|
}
|
|
connUUID, ok := common.CheckDevice(base.Device, base.Conn)
|
|
if !ok {
|
|
ctx.AbortWithStatusJSON(http.StatusBadGateway, modules.Packet{Code: 1, Msg: `${i18n|deviceNotExists}`})
|
|
return ``, false
|
|
}
|
|
ctx.Request = ctx.Request.WithContext(context.WithValue(ctx.Request.Context(), `ConnUUID`, connUUID))
|
|
return connUUID, true
|
|
}
|
|
|
|
// OnDevicePack handles events about device info.
|
|
// Such as websocket handshake and update device info.
|
|
func OnDevicePack(data []byte, session *melody.Session) error {
|
|
var pack struct {
|
|
Code int `json:"code,omitempty"`
|
|
Act string `json:"act,omitempty"`
|
|
Msg string `json:"msg,omitempty"`
|
|
Device modules.Device `json:"data"`
|
|
}
|
|
err := utils.JSON.Unmarshal(data, &pack)
|
|
if err != nil {
|
|
golog.Error(err)
|
|
session.Close()
|
|
return err
|
|
}
|
|
|
|
addr, ok := session.Get(`Address`)
|
|
if ok {
|
|
pack.Device.WAN = addr.(string)
|
|
} else {
|
|
pack.Device.WAN = `Unknown`
|
|
}
|
|
|
|
if pack.Act == `report` {
|
|
// Check if this device has already connected.
|
|
// If so, then find the session and let client quit.
|
|
// This will keep only one connection remained per device.
|
|
exSession := ``
|
|
common.Devices.IterCb(func(uuid string, v any) bool {
|
|
device := v.(*modules.Device)
|
|
if device.ID == pack.Device.ID {
|
|
exSession = uuid
|
|
target, ok := common.Melody.GetSessionByUUID(uuid)
|
|
if ok {
|
|
common.SendPack(modules.Packet{Act: `offline`}, target)
|
|
target.Close()
|
|
}
|
|
return false
|
|
}
|
|
return true
|
|
})
|
|
if len(exSession) > 0 {
|
|
common.Devices.Remove(exSession)
|
|
}
|
|
common.Devices.Set(session.UUID, &pack.Device)
|
|
common.Info(nil, `CLIENT_ONLINE`, ``, ``, map[string]any{
|
|
`device`: map[string]any{
|
|
`name`: pack.Device.Hostname,
|
|
`ip`: pack.Device.WAN,
|
|
},
|
|
})
|
|
} else {
|
|
val, ok := common.Devices.Get(session.UUID)
|
|
if ok {
|
|
deviceInfo := val.(*modules.Device)
|
|
deviceInfo.CPU = pack.Device.CPU
|
|
deviceInfo.RAM = pack.Device.RAM
|
|
deviceInfo.Net = pack.Device.Net
|
|
deviceInfo.Disk = pack.Device.Disk
|
|
deviceInfo.Uptime = pack.Device.Uptime
|
|
}
|
|
}
|
|
common.SendPack(modules.Packet{Code: 0}, session)
|
|
return nil
|
|
}
|
|
|
|
// CheckUpdate will check if client need update and return latest client if so.
|
|
func CheckUpdate(ctx *gin.Context) {
|
|
var form struct {
|
|
OS string `form:"os" binding:"required"`
|
|
Arch string `form:"arch" binding:"required"`
|
|
Commit string `form:"commit" binding:"required"`
|
|
}
|
|
if err := ctx.ShouldBind(&form); err != nil {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return
|
|
}
|
|
if form.Commit == config.COMMIT {
|
|
ctx.JSON(http.StatusOK, modules.Packet{Code: 0})
|
|
common.Warn(ctx, `CLIENT_UPDATE`, `success`, `latest`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
return
|
|
}
|
|
tpl, err := os.Open(fmt.Sprintf(config.BuiltPath, form.OS, form.Arch))
|
|
if err != nil {
|
|
ctx.AbortWithStatusJSON(http.StatusNotFound, modules.Packet{Code: 1, Msg: `${i18n|osOrArchNotPrebuilt}`})
|
|
common.Warn(ctx, `CLIENT_UPDATE`, `fail`, `no prebuild asset`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
return
|
|
}
|
|
|
|
const MaxBodySize = 384 // This is size of client config buffer.
|
|
if ctx.Request.ContentLength > MaxBodySize {
|
|
ctx.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, modules.Packet{Code: 1})
|
|
common.Warn(ctx, `CLIENT_UPDATE`, `fail`, `config too large`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
return
|
|
}
|
|
body, err := ctx.GetRawData()
|
|
if err != nil {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: 1})
|
|
common.Warn(ctx, `CLIENT_UPDATE`, `fail`, `read config fail`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
return
|
|
}
|
|
session := common.CheckClientReq(ctx)
|
|
if session == nil {
|
|
ctx.AbortWithStatusJSON(http.StatusUnauthorized, modules.Packet{Code: 1})
|
|
common.Warn(ctx, `CLIENT_UPDATE`, `fail`, `check config fail`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
return
|
|
}
|
|
|
|
common.Info(ctx, `CLIENT_UPDATE`, `success`, `updating`, map[string]any{
|
|
`client`: map[string]any{
|
|
`os`: form.OS,
|
|
`arch`: form.Arch,
|
|
`commit`: form.Commit,
|
|
},
|
|
`server`: config.COMMIT,
|
|
})
|
|
|
|
ctx.Header(`Spark-Commit`, config.COMMIT)
|
|
ctx.Header(`Accept-Ranges`, `none`)
|
|
ctx.Header(`Content-Transfer-Encoding`, `binary`)
|
|
ctx.Header(`Content-Type`, `application/octet-stream`)
|
|
if stat, err := tpl.Stat(); err == nil {
|
|
ctx.Header(`Content-Length`, strconv.FormatInt(stat.Size(), 10))
|
|
}
|
|
cfgBuffer := bytes.Repeat([]byte{'\x19'}, 384)
|
|
prevBuffer := make([]byte, 0)
|
|
for {
|
|
thisBuffer := make([]byte, 1024)
|
|
n, err := tpl.Read(thisBuffer)
|
|
thisBuffer = thisBuffer[:n]
|
|
tempBuffer := append(prevBuffer, thisBuffer...)
|
|
bufIndex := bytes.Index(tempBuffer, cfgBuffer)
|
|
if bufIndex > -1 {
|
|
tempBuffer = bytes.Replace(tempBuffer, cfgBuffer, body, -1)
|
|
}
|
|
ctx.Writer.Write(tempBuffer[:len(prevBuffer)])
|
|
prevBuffer = tempBuffer[len(prevBuffer):]
|
|
if err != nil {
|
|
break
|
|
}
|
|
}
|
|
if len(prevBuffer) > 0 {
|
|
ctx.Writer.Write(prevBuffer)
|
|
prevBuffer = []byte{}
|
|
}
|
|
}
|
|
|
|
// ExecDeviceCmd execute command on device.
|
|
func ExecDeviceCmd(ctx *gin.Context) {
|
|
var form struct {
|
|
Cmd string `json:"cmd" yaml:"cmd" form:"cmd" binding:"required"`
|
|
Args string `json:"args" yaml:"args" form:"args"`
|
|
}
|
|
target, ok := CheckForm(ctx, &form)
|
|
if !ok {
|
|
return
|
|
}
|
|
if len(form.Cmd) == 0 {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return
|
|
}
|
|
trigger := utils.GetStrUUID()
|
|
common.SendPackByUUID(modules.Packet{Code: 0, Act: `execCommand`, Data: gin.H{`cmd`: form.Cmd, `args`: form.Args}, Event: trigger}, target)
|
|
ok = common.AddEventOnce(func(p modules.Packet, _ *melody.Session) {
|
|
if p.Code != 0 {
|
|
common.Warn(ctx, `EXEC_COMMAND`, `fail`, p.Msg, map[string]any{
|
|
`cmd`: form.Cmd,
|
|
`args`: form.Args,
|
|
})
|
|
ctx.AbortWithStatusJSON(http.StatusInternalServerError, modules.Packet{Code: 1, Msg: p.Msg})
|
|
} else {
|
|
common.Info(ctx, `EXEC_COMMAND`, `success`, ``, map[string]any{
|
|
`cmd`: form.Cmd,
|
|
`args`: form.Args,
|
|
})
|
|
ctx.JSON(http.StatusOK, modules.Packet{Code: 0})
|
|
}
|
|
}, target, trigger, 5*time.Second)
|
|
if !ok {
|
|
common.Warn(ctx, `EXEC_COMMAND`, `fail`, `timeout`, map[string]any{
|
|
`cmd`: form.Cmd,
|
|
`args`: form.Args,
|
|
})
|
|
ctx.AbortWithStatusJSON(http.StatusGatewayTimeout, modules.Packet{Code: 1, Msg: `${i18n|responseTimeout}`})
|
|
}
|
|
}
|
|
|
|
// GetDevices will return all info about all clients.
|
|
func GetDevices(ctx *gin.Context) {
|
|
devices := map[string]any{}
|
|
common.Devices.IterCb(func(uuid string, v any) bool {
|
|
device := v.(*modules.Device)
|
|
devices[uuid] = *device
|
|
return true
|
|
})
|
|
ctx.JSON(http.StatusOK, modules.Packet{Code: 0, Data: devices})
|
|
}
|
|
|
|
// CallDevice will call client with command from browser.
|
|
func CallDevice(ctx *gin.Context) {
|
|
act := ctx.Param(`act`)
|
|
if len(act) == 0 {
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return
|
|
}
|
|
{
|
|
actions := []string{`lock`, `logoff`, `hibernate`, `suspend`, `restart`, `shutdown`, `offline`}
|
|
ok := false
|
|
for _, v := range actions {
|
|
if v == act {
|
|
ok = true
|
|
break
|
|
}
|
|
}
|
|
if !ok {
|
|
common.Warn(ctx, `CALL_DEVICE`, `fail`, `invalid act`, map[string]any{
|
|
`act`: act,
|
|
})
|
|
ctx.AbortWithStatusJSON(http.StatusBadRequest, modules.Packet{Code: -1, Msg: `${i18n|invalidParameter}`})
|
|
return
|
|
}
|
|
}
|
|
connUUID, ok := CheckForm(ctx, nil)
|
|
if !ok {
|
|
return
|
|
}
|
|
trigger := utils.GetStrUUID()
|
|
common.SendPackByUUID(modules.Packet{Act: act, Event: trigger}, connUUID)
|
|
ok = common.AddEventOnce(func(p modules.Packet, _ *melody.Session) {
|
|
if p.Code != 0 {
|
|
common.Warn(ctx, `CALL_DEVICE`, `fail`, p.Msg, map[string]any{
|
|
`act`: act,
|
|
})
|
|
ctx.AbortWithStatusJSON(http.StatusInternalServerError, modules.Packet{Code: 1, Msg: p.Msg})
|
|
} else {
|
|
common.Info(ctx, `CALL_DEVICE`, `success`, ``, map[string]any{
|
|
`act`: act,
|
|
})
|
|
ctx.JSON(http.StatusOK, modules.Packet{Code: 0})
|
|
}
|
|
}, connUUID, trigger, 5*time.Second)
|
|
if !ok {
|
|
//This means the client is offline.
|
|
//So we take this as a success.
|
|
common.Info(ctx, `CALL_DEVICE`, `success`, ``, map[string]any{
|
|
`act`: act,
|
|
})
|
|
ctx.JSON(http.StatusOK, modules.Packet{Code: 0})
|
|
}
|
|
}
|
|
|
|
func SimpleEncrypt(data []byte, session *melody.Session) ([]byte, bool) {
|
|
temp, ok := session.Get(`Secret`)
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
secret := temp.([]byte)
|
|
block, err := aes.NewCipher(secret)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
stream := cipher.NewCTR(block, secret)
|
|
encBuffer := make([]byte, len(data))
|
|
stream.XORKeyStream(encBuffer, data)
|
|
return encBuffer, true
|
|
}
|
|
|
|
func SimpleDecrypt(data []byte, session *melody.Session) ([]byte, bool) {
|
|
temp, ok := session.Get(`Secret`)
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
secret := temp.([]byte)
|
|
block, err := aes.NewCipher(secret)
|
|
if err != nil {
|
|
return nil, false
|
|
}
|
|
stream := cipher.NewCTR(block, secret)
|
|
decBuffer := make([]byte, len(data))
|
|
stream.XORKeyStream(decBuffer, data)
|
|
return decBuffer, true
|
|
}
|
|
|
|
func WSHealthCheck(container *melody.Melody, sender Sender) {
|
|
const MaxIdleSeconds = 300
|
|
ping := func(uuid string, s *melody.Session) {
|
|
if !sender(modules.Packet{Act: `ping`}, s) {
|
|
s.Close()
|
|
}
|
|
}
|
|
for now := range time.NewTicker(60 * time.Second).C {
|
|
timestamp := now.Unix()
|
|
// stores sessions to be disconnected
|
|
queue := make([]*melody.Session, 0)
|
|
container.IterSessions(func(uuid string, s *melody.Session) bool {
|
|
go ping(uuid, s)
|
|
val, ok := s.Get(`LastPack`)
|
|
if !ok {
|
|
queue = append(queue, s)
|
|
return true
|
|
}
|
|
lastPack, ok := val.(int64)
|
|
if !ok {
|
|
queue = append(queue, s)
|
|
return true
|
|
}
|
|
if timestamp-lastPack > MaxIdleSeconds {
|
|
queue = append(queue, s)
|
|
}
|
|
return true
|
|
})
|
|
for i := 0; i < len(queue); i++ {
|
|
queue[i].Close()
|
|
}
|
|
}
|
|
}
|