gRPC 上报系统信息

service/handler/auth.go

@@ -0,0 +1,60 @@
+package handler
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+// AuthHandler ..
+type AuthHandler struct {
+	AppKey    string
+	AppSecret string
+}
+
+// GetRequestMetadata ..
+func (a *AuthHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+	return map[string]string{"app_key": a.AppKey, "app_secret": a.AppSecret}, nil
+}
+
+// RequireTransportSecurity ..
+func (a *AuthHandler) RequireTransportSecurity() bool {
+	return false
+}
+
+// Check ..
+func (a *AuthHandler) Check(ctx context.Context) error {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return status.Errorf(codes.Unauthenticated, "metadata.FromIncomingContext err")
+	}
+
+	var (
+		AppKey    string
+		AppSecret string
+	)
+	if value, ok := md["app_key"]; ok {
+		AppKey = value[0]
+	}
+	if value, ok := md["app_secret"]; ok {
+		AppSecret = value[0]
+	}
+
+	if AppKey != a.GetAppKey() || AppSecret != a.GetAppSecret() {
+		return status.Errorf(codes.Unauthenticated, "invalid token")
+	}
+
+	return nil
+}
+
+// GetAppKey ..
+func (a *AuthHandler) GetAppKey() string {
+	return a.AppKey
+}
+
+// GetAppSecret ..
+func (a *AuthHandler) GetAppSecret() string {
+	return a.AppSecret
+}

service/handler/nezha.go

@@ -0,0 +1,40 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+
+	pb "github.com/p14yground/nezha/proto"
+)
+
+// NezhaHandler ..
+type NezhaHandler struct {
+	Auth *AuthHandler
+}
+
+// ReportState ..
+func (s *NezhaHandler) ReportState(c context.Context, r *pb.State) (*pb.Receipt, error) {
+	if err := s.Auth.Check(c); err != nil {
+		return nil, err
+	}
+	fmt.Printf("ReportState receive: %s\n", r)
+	return &pb.Receipt{Proced: true}, nil
+}
+
+// Heartbeat ..
+func (s *NezhaHandler) Heartbeat(r *pb.Beat, stream pb.NezhaService_HeartbeatServer) error {
+	if err := s.Auth.Check(stream.Context()); err != nil {
+		return err
+	}
+	fmt.Printf("ReportState receive: %s\n", r)
+	return nil
+}
+
+// Register ..
+func (s *NezhaHandler) Register(c context.Context, r *pb.Host) (*pb.Receipt, error) {
+	if err := s.Auth.Check(c); err != nil {
+		return nil, err
+	}
+	fmt.Printf("Register receive: %s\n", r)
+	return &pb.Receipt{Proced: true}, nil
+}

service/monitor/monitor.go

@@ -0,0 +1,73 @@
+package monitor
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/shirou/gopsutil/cpu"
+	"github.com/shirou/gopsutil/disk"
+	"github.com/shirou/gopsutil/host"
+	"github.com/shirou/gopsutil/mem"
+	"github.com/shirou/gopsutil/net"
+
+	"github.com/p14yground/nezha/model"
+)
+
+// GetHost ..
+func GetHost() *model.Host {
+	hi, _ := host.Info()
+	var cpus []string
+	ci, _ := cpu.Info()
+	for i := 0; i < len(ci); i++ {
+		cpus = append(cpus, fmt.Sprintf("%v-%vC%vT", ci[i].ModelName, ci[i].Cores, ci[i].Stepping))
+	}
+	return &model.Host{
+		Platform:        hi.OS,
+		PlatformVersion: hi.PlatformVersion,
+		CPU:             cpus,
+		Arch:            hi.KernelArch,
+		Virtualization:  hi.VirtualizationSystem,
+		Uptime:          fmt.Sprintf("%v", hi.Uptime),
+		BootTime:        fmt.Sprintf("%v", hi.BootTime),
+	}
+}
+
+// GetState ..
+func GetState(delay uint64) *model.State {
+	// Memory
+	mv, _ := mem.VirtualMemory()
+	ms, _ := mem.SwapMemory()
+	// Disk
+	var diskTotal, diskUsed uint64
+	dparts, _ := disk.Partitions(true)
+	for _, part := range dparts {
+		u, _ := disk.Usage(part.Mountpoint)
+		diskTotal += u.Total
+		diskUsed += u.Used
+	}
+	// CPU
+	var cpuPercent float64
+	cp, err := cpu.Percent(time.Second*time.Duration(delay), false)
+	if err == nil {
+		cpuPercent = cp[0]
+
+	}
+	// Network
+	var netIn, netOut uint64
+	nc, err := net.IOCounters(false)
+	if err == nil {
+		netIn = nc[0].BytesRecv
+		netOut = nc[0].BytesSent
+	}
+	return &model.State{
+		CPU:       cpuPercent,
+		MEMTotal:  mv.Total,
+		MEMUsed:   mv.Used,
+		SwapTotal: ms.Total,
+		SwapUsed:  ms.Used,
+		DiskTotal: diskTotal,
+		DiskUsed:  diskUsed,
+		NetIn:     netIn,
+		NetOut:    netOut,
+	}
+}