diff --git a/args/args.go b/args/args.go index d6a637ae..df20932f 100644 --- a/args/args.go +++ b/args/args.go @@ -22,9 +22,12 @@ type BotRoom struct { } var Bot = struct { - Enable bool `yaml:"enable"` - Welcome string `yaml:"welcome"` - InvitableRooms []*BotRoom `yaml:"invitableRooms"` + Enable bool `yaml:"enable"` + Welcome string `yaml:"welcome"` + Managers []string `yaml:"managers"` + BlackList []string `yaml:"blackList"` + WhiteList []string `yaml:"whiteList"` + HostedRooms []*BotRoom `yaml:"hostedRooms"` }{ Enable: true, } diff --git a/config.yml b/config.yml index b3e2dcb8..2a042cf9 100644 --- a/config.yml +++ b/config.yml @@ -2,7 +2,10 @@ bot: enable: true # 是否启用内置机器人 welcome: 回复 /help 查询指令 # 添加好友或加群后自动回复 - invitableRooms: # 可邀请的群组列表 + managers: [] # 管理员人员 + blackList: [] # 黑名单,禁止使用机器人 + whiteList: [] # 白名单,禁止其他人使用机器人 + hostedRooms: # 托管的群聊列表 - mask: 1 # 群标记 name: OpenTDP 开发群 # 群聊名称 roomId: 38699745819@chatroom # 群聊 Id diff --git a/wclient/proto/friend.go b/wclient/proto/msg_friend.go similarity index 100% rename from wclient/proto/friend.go rename to wclient/proto/msg_friend.go diff --git a/wclient/proto/xml_atuser.go b/wclient/proto/xml_atuser.go new file mode 100644 index 00000000..e54803fd --- /dev/null +++ b/wclient/proto/xml_atuser.go @@ -0,0 +1,13 @@ +package proto + +type AtMsgSource struct { + AtUserList string `xml:"atuserlist"` + Silence int32 `xml:"silence"` + MemberCount int32 `xml:"membercount"` + Signature string `xml:"signature"` + TmpNode TmpNode `xml:"tmp_node"` +} + +type TmpNode struct { + PublisherID string `xml:",chardata"` +} diff --git a/wclient/robot/handle.go b/wclient/robot/handle.go index 99753035..fd2eb82d 100644 --- a/wclient/robot/handle.go +++ b/wclient/robot/handle.go @@ -1,6 +1,7 @@ package robot import ( + "encoding/xml" "fmt" "math/rand" "regexp" @@ -10,6 +11,7 @@ import ( "github.com/opentdp/wechat-rest/args" "github.com/opentdp/wechat-rest/wcferry" "github.com/opentdp/wechat-rest/wclient/model" + "github.com/opentdp/wechat-rest/wclient/proto" ) type Handler struct { @@ -47,6 +49,27 @@ func initHandlers() { }, } + handlers["/ban"] = &Handler{ + Level: 1, + ChatAble: true, + RoomAble: true, + Describe: "禁止用户使用Ai服务", + Callback: func(msg *wcferry.WxMsg) string { + ret := &proto.AtMsgSource{} + err := xml.Unmarshal([]byte(msg.Xml), ret) + if err == nil && ret.AtUserList != "" { + users := strings.Split(ret.AtUserList, ",") + for _, v := range users { + if v != "" && !contains(args.Bot.BlackList, v) { + args.Bot.BlackList = append(args.Bot.BlackList, v) + } + } + return fmt.Sprintf("操作完成,已禁止用户数:%d", len(args.Bot.BlackList)) + } + return "参数错误" + }, + } + handlers["/mr"] = &Handler{ Level: 0, ChatAble: true, @@ -72,7 +95,7 @@ func initHandlers() { } } - for _, v := range args.Bot.InvitableRooms { + for _, v := range args.Bot.HostedRooms { v := v // copy it cmdkey := "/room:" + v.Mask handlers[cmdkey] = &Handler{ @@ -127,6 +150,14 @@ func initHandlers() { func applyHandler(msg *wcferry.WxMsg) string { + // 匹配名单 + if len(args.Bot.WhiteList) > 0 && !contains(args.Bot.WhiteList, msg.Sender) { + return "" + } + if len(args.Bot.BlackList) > 0 && contains(args.Bot.BlackList, msg.Sender) { + return "" + } + // 解析指令 re := regexp.MustCompile(`^(/[\w:-]{2,20})(\s|$)`) matches := re.FindStringSubmatch(msg.Content) @@ -148,7 +179,7 @@ func applyHandler(msg *wcferry.WxMsg) string { } // 检查权限 - if hd.Level > 0 { + if hd.Level > 0 && !contains(args.Bot.Managers, msg.Sender) { return "无权限使用该指令" } @@ -167,3 +198,12 @@ func applyHandler(msg *wcferry.WxMsg) string { return hd.Callback(msg) } + +func contains(slice []string, val string) bool { + for _, item := range slice { + if item == val { + return true + } + } + return false +}