From e28e7f4375b973a66d6c88be04b45f5661b9be78 Mon Sep 17 00:00:00 2001
From: Lemoe <liukangdev@gmail.com>
Date: Sat, 6 Nov 2021 09:46:44 +0800
Subject: [PATCH] =?UTF-8?q?agent=20=E5=A2=9E=E5=8A=A0=20SSL/TLS=20?=
 =?UTF-8?q?=E9=80=89=E9=A1=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/agent/main.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 71510c5..ffe5c1f 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -18,6 +19,7 @@ import (
 	"github.com/p14yground/go-github-selfupdate/selfupdate"
 	flag "github.com/spf13/pflag"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
 
 	"github.com/naiba/nezha/cmd/agent/monitor"
 	"github.com/naiba/nezha/cmd/agent/processgroup"
@@ -38,6 +40,7 @@ type AgentConfig struct {
 	Server                string
 	ClientSecret          string
 	ReportDelay           int
+	TLS                   bool
 }
 
 var (
@@ -80,6 +83,7 @@ func main() {
 	flag.BoolVar(&agentConf.DisableCommandExecute, "disable-command-execute", false, "禁止在此机器上执行命令")
 	flag.BoolVar(&agentConf.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级")
 	flag.BoolVar(&agentConf.DisableForceUpdate, "disable-force-update", false, "禁用强制升级")
+	flag.BoolVar(&agentConf.TLS, "tls", false, "启用SSL/TLS加密")
 	flag.Parse()
 
 	if agentConf.ClientSecret == "" {
@@ -138,7 +142,13 @@ func run() {
 
 	for {
 		timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
-		conn, err = grpc.DialContext(timeOutCtx, agentConf.Server, grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
+		var securityOption grpc.DialOption
+		if agentConf.TLS {
+			securityOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{}))
+		} else {
+			securityOption = grpc.WithInsecure()
+		}
+		conn, err = grpc.DialContext(timeOutCtx, agentConf.Server, securityOption, grpc.WithPerRPCCredentials(&auth))
 		if err != nil {
 			println("与面板建立连接失败：", err)
 			cancel()