mirror of
https://github.com/TeaOSLab/EdgeNode.git
synced 2025-11-04 07:40:56 +08:00
341 lines
8.2 KiB
Go
341 lines
8.2 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"github.com/TeaOSLab/EdgeNode/internal/apps"
|
|
teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
|
|
"github.com/TeaOSLab/EdgeNode/internal/nodes"
|
|
_ "github.com/iwind/TeaGo/bootstrap"
|
|
"github.com/iwind/TeaGo/logs"
|
|
"github.com/iwind/TeaGo/maps"
|
|
"github.com/iwind/TeaGo/types"
|
|
"github.com/iwind/gosock/pkg/gosock"
|
|
"net"
|
|
"net/http"
|
|
_ "net/http/pprof"
|
|
"os"
|
|
"sort"
|
|
)
|
|
|
|
func main() {
|
|
var app = apps.NewAppCmd().
|
|
Version(teaconst.Version).
|
|
Product(teaconst.ProductName).
|
|
Usage(teaconst.ProcessName + " [-v|start|stop|restart|status|quit|test|reload|service|daemon|pprof|accesslog]").
|
|
Usage(teaconst.ProcessName + " [trackers|goman|conns|gc]").
|
|
Usage(teaconst.ProcessName + " [ip.drop|ip.reject|ip.remove] IP")
|
|
|
|
app.On("test", func() {
|
|
err := nodes.NewNode().Test()
|
|
if err != nil {
|
|
_, _ = os.Stderr.WriteString(err.Error())
|
|
}
|
|
})
|
|
app.On("reload", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{Code: "reload"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
var params = maps.NewMap(reply.Params)
|
|
if params.Has("error") {
|
|
fmt.Println("[ERROR]" + params.GetString("error"))
|
|
} else {
|
|
fmt.Println("ok")
|
|
}
|
|
}
|
|
})
|
|
app.On("daemon", func() {
|
|
nodes.NewNode().Daemon()
|
|
})
|
|
app.On("service", func() {
|
|
err := nodes.NewNode().InstallSystemService()
|
|
if err != nil {
|
|
fmt.Println("[ERROR]install failed: " + err.Error())
|
|
return
|
|
}
|
|
fmt.Println("done")
|
|
})
|
|
app.On("quit", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
_, err := sock.Send(&gosock.Command{Code: "quit"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]quit failed: " + err.Error())
|
|
return
|
|
}
|
|
fmt.Println("done")
|
|
})
|
|
app.On("pprof", func() {
|
|
var flagSet = flag.NewFlagSet("pprof", flag.ExitOnError)
|
|
var addr string
|
|
flagSet.StringVar(&addr, "addr", "", "")
|
|
_ = flagSet.Parse(os.Args[2:])
|
|
|
|
if len(addr) == 0 {
|
|
addr = "127.0.0.1:6060"
|
|
}
|
|
logs.Println("starting with pprof '" + addr + "'...")
|
|
|
|
go func() {
|
|
err := http.ListenAndServe(addr, nil)
|
|
if err != nil {
|
|
logs.Println("[ERROR]" + err.Error())
|
|
}
|
|
}()
|
|
|
|
var node = nodes.NewNode()
|
|
node.Start()
|
|
})
|
|
app.On("dbstat", func() {
|
|
teaconst.EnableDBStat = true
|
|
|
|
var node = nodes.NewNode()
|
|
node.Start()
|
|
})
|
|
app.On("trackers", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{Code: "trackers"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
labelsMap, ok := reply.Params["labels"]
|
|
if ok {
|
|
labels, ok := labelsMap.(map[string]interface{})
|
|
if ok {
|
|
if len(labels) == 0 {
|
|
fmt.Println("no labels yet")
|
|
} else {
|
|
var labelNames = []string{}
|
|
for label := range labels {
|
|
labelNames = append(labelNames, label)
|
|
}
|
|
sort.Strings(labelNames)
|
|
|
|
for _, labelName := range labelNames {
|
|
fmt.Println(labelName + ": " + fmt.Sprintf("%.6f", labels[labelName]))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
})
|
|
app.On("goman", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{Code: "goman"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
instancesJSON, err := json.MarshalIndent(reply.Params, "", " ")
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
fmt.Println(string(instancesJSON))
|
|
}
|
|
}
|
|
})
|
|
app.On("conns", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{Code: "conns"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
resultJSON, err := json.MarshalIndent(reply.Params, "", " ")
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
fmt.Println(string(resultJSON))
|
|
}
|
|
}
|
|
})
|
|
app.On("gc", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
_, err := sock.Send(&gosock.Command{Code: "gc"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
fmt.Println("ok")
|
|
}
|
|
})
|
|
app.On("ip.drop", func() {
|
|
var args = os.Args[2:]
|
|
if len(args) == 0 {
|
|
fmt.Println("Usage: edge-node ip.drop IP [--timeout=SECONDS] [--async]")
|
|
return
|
|
}
|
|
var ip = args[0]
|
|
if len(net.ParseIP(ip)) == 0 {
|
|
fmt.Println("IP '" + ip + "' is invalid")
|
|
return
|
|
}
|
|
var timeoutSeconds = 0
|
|
var options = app.ParseOptions(args[1:])
|
|
timeout, ok := options["timeout"]
|
|
if ok {
|
|
timeoutSeconds = types.Int(timeout[0])
|
|
}
|
|
var async = false
|
|
_, ok = options["async"]
|
|
if ok {
|
|
async = true
|
|
}
|
|
|
|
fmt.Println("drop ip '" + ip + "' for '" + types.String(timeoutSeconds) + "' seconds")
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{
|
|
Code: "dropIP",
|
|
Params: map[string]interface{}{
|
|
"ip": ip,
|
|
"timeoutSeconds": timeoutSeconds,
|
|
"async": async,
|
|
},
|
|
})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
var errString = maps.NewMap(reply.Params).GetString("error")
|
|
if len(errString) > 0 {
|
|
fmt.Println("[ERROR]" + errString)
|
|
} else {
|
|
fmt.Println("ok")
|
|
}
|
|
}
|
|
})
|
|
app.On("ip.reject", func() {
|
|
var args = os.Args[2:]
|
|
if len(args) == 0 {
|
|
fmt.Println("Usage: edge-node ip.reject IP [--timeout=SECONDS]")
|
|
return
|
|
}
|
|
var ip = args[0]
|
|
if len(net.ParseIP(ip)) == 0 {
|
|
fmt.Println("IP '" + ip + "' is invalid")
|
|
return
|
|
}
|
|
var timeoutSeconds = 0
|
|
var options = app.ParseOptions(args[1:])
|
|
timeout, ok := options["timeout"]
|
|
if ok {
|
|
timeoutSeconds = types.Int(timeout[0])
|
|
}
|
|
|
|
fmt.Println("reject ip '" + ip + "' for '" + types.String(timeoutSeconds) + "' seconds")
|
|
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{
|
|
Code: "rejectIP",
|
|
Params: map[string]interface{}{
|
|
"ip": ip,
|
|
"timeoutSeconds": timeoutSeconds,
|
|
},
|
|
})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
var errString = maps.NewMap(reply.Params).GetString("error")
|
|
if len(errString) > 0 {
|
|
fmt.Println("[ERROR]" + errString)
|
|
} else {
|
|
fmt.Println("ok")
|
|
}
|
|
}
|
|
})
|
|
app.On("ip.remove", func() {
|
|
var args = os.Args[2:]
|
|
if len(args) == 0 {
|
|
fmt.Println("Usage: edge-node ip.remove IP")
|
|
return
|
|
}
|
|
var ip = args[0]
|
|
if len(net.ParseIP(ip)) == 0 {
|
|
fmt.Println("IP '" + ip + "' is invalid")
|
|
return
|
|
}
|
|
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{
|
|
Code: "removeIP",
|
|
Params: map[string]interface{}{
|
|
"ip": ip,
|
|
},
|
|
})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
} else {
|
|
var errString = maps.NewMap(reply.Params).GetString("error")
|
|
if len(errString) > 0 {
|
|
fmt.Println("[ERROR]" + errString)
|
|
} else {
|
|
fmt.Println("ok")
|
|
}
|
|
}
|
|
})
|
|
app.On("accesslog", func() {
|
|
// local sock
|
|
var tmpDir = os.TempDir()
|
|
var sockFile = tmpDir + "/" + teaconst.AccessLogSockName
|
|
_, err := os.Stat(sockFile)
|
|
if err != nil {
|
|
if !os.IsNotExist(err) {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
return
|
|
}
|
|
}
|
|
|
|
var processSock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := processSock.Send(&gosock.Command{
|
|
Code: "accesslog",
|
|
})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
return
|
|
}
|
|
if reply.Code == "error" {
|
|
var errString = maps.NewMap(reply.Params).GetString("error")
|
|
if len(errString) > 0 {
|
|
fmt.Println("[ERROR]" + errString)
|
|
return
|
|
}
|
|
}
|
|
|
|
conn, err := net.Dial("unix", sockFile)
|
|
if err != nil {
|
|
fmt.Println("[ERROR]start reading access log failed: " + err.Error())
|
|
return
|
|
}
|
|
defer func() {
|
|
_ = conn.Close()
|
|
}()
|
|
var buf = make([]byte, 1024)
|
|
for {
|
|
n, err := conn.Read(buf)
|
|
if n > 0 {
|
|
fmt.Print(string(buf[:n]))
|
|
}
|
|
if err != nil {
|
|
break
|
|
}
|
|
}
|
|
})
|
|
app.On("bandwidth", func() {
|
|
var sock = gosock.NewTmpSock(teaconst.ProcessName)
|
|
reply, err := sock.Send(&gosock.Command{Code: "bandwidth"})
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
return
|
|
}
|
|
var statsMap = maps.NewMap(reply.Params).Get("stats")
|
|
statsJSON, err := json.MarshalIndent(statsMap, "", " ")
|
|
if err != nil {
|
|
fmt.Println("[ERROR]" + err.Error())
|
|
return
|
|
}
|
|
fmt.Println(string(statsJSON))
|
|
})
|
|
app.Run(func() {
|
|
var node = nodes.NewNode()
|
|
node.Start()
|
|
})
|
|
}
|