添加清空流量数据命令

cmd/clean.go

@@ -0,0 +1,23 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+	"trojan/trojan"
+)
+
+// cleanCmd represents the clean command
+var cleanCmd = &cobra.Command{
+	Use:   "clean",
+	Short: "清空指定用户流量",
+	Long: `传入指定用户名来清空用户流量, 多个空格隔开, 例如:
+trojan clean zhangsan lisi
+`,
+	Args: cobra.MinimumNArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		trojan.CleanDataByName(args)
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(cleanCmd)
+}

core/mysql.go

@@ -6,15 +6,17 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	mysqlDriver "github.com/go-sql-driver/mysql"
 	"io/ioutil"
 	"log"
 	"trojan/util"
 
+	mysqlDriver "github.com/go-sql-driver/mysql"
+
 	// mysql sql驱动
-	_ "github.com/go-sql-driver/mysql"
 	"strconv"
 	"strings"
+
+	_ "github.com/go-sql-driver/mysql"
 )
 
 // Mysql 结构体
@@ -192,6 +194,29 @@ func (mysql *Mysql) CleanData(id uint) error {
 	return nil
 }
 
+// CleanDataByName 清空指定用户名流量统计数据
+func (mysql *Mysql) CleanDataByName(usernames []string) error {
+	db := mysql.GetDB()
+	if db == nil {
+		return errors.New("can't connect mysql")
+	}
+	defer db.Close()
+	runSql := "UPDATE users SET download=0, upload=0 WHERE username in ("
+	for i, name := range usernames {
+		runSql = runSql + "'" + name + "'"
+		if i == len(usernames)-1 {
+			runSql = runSql + ")"
+		} else {
+			runSql = runSql + ","
+		}
+	}
+	if _, err := db.Exec(runSql); err != nil {
+		fmt.Println(err)
+		return err
+	}
+	return nil
+}
+
 // GetUserByName 通过用户名来获取用户
 func (mysql *Mysql) GetUserByName(name string) *User {
 	db := mysql.GetDB()

trojan/user.go

@@ -101,6 +101,14 @@ func CleanData() {
 	}
 }
 
+// CleanDataByName 清空指定用户流量
+func CleanDataByName(usernames []string) {
+	mysql := core.GetMysql()
+	if err := mysql.CleanDataByName(usernames); err != nil {
+		fmt.Println(err.Error())
+	}
+}
+
 // UserList 获取用户列表并打印显示
 func UserList(ids ...string) []*core.User {
 	mysql := core.GetMysql()