Filter available commands by chat type

3 files changed, 143 insertions, 55 deletions

diff --git a/telegram/commands.go b/telegram/commands.go
index 8d4de91..a01a80e 100644
--- a/telegram/commands.go
+++ b/telegram/commands.go
@@ -50,56 +50,60 @@ var permissionsMember = client.ChatPermissions{
 var permissionsReadonly = client.ChatPermissions{}
 
 var transportCommands = map[string]command{
-	"help":        command{0, []string{}, "help"},
-	"login":       command{1, []string{"phone"}, "sign in"},
-	"logout":      command{0, []string{}, "sign out"},
-	"cancelauth":  command{0, []string{}, "quit the signin wizard"},
-	"code":        command{1, []string{"xxxxx"}, "check one-time code"},
-	"password":    command{1, []string{"********"}, "check 2fa password"},
-	"setusername": command{0, []string{"@username"}, "update @username"},
-	"setname":     command{1, []string{"first", "last"}, "update name"},
-	"setbio":      command{0, []string{"Lorem ipsum"}, "update about"},
-	"setpassword": command{0, []string{"old", "new"}, "set or remove password"},
-	"config":      command{0, []string{"param", "value"}, "view or update configuration options"},
-	"report":      command{2, []string{"chat", "comment"}, "report a chat by id or @username"},
-	"add":         command{1, []string{"@username"}, "add @username to your chat list"},
-	"join":        command{1, []string{"https://t.me/invite_link"}, "join to chat via invite link or @publicname"},
-	"supergroup":  command{1, []string{"title", "description"}, "create new supergroup «title» with «description»"},
-	"channel":     command{1, []string{"title", "description"}, "create new channel «title» with «description»"},
+	"help":        command{0, []string{}, "help", nil},
+	"login":       command{1, []string{"phone"}, "sign in", nil},
+	"logout":      command{0, []string{}, "sign out", nil},
+	"cancelauth":  command{0, []string{}, "quit the signin wizard", nil},
+	"code":        command{1, []string{"xxxxx"}, "check one-time code", nil},
+	"password":    command{1, []string{"********"}, "check 2fa password", nil},
+	"setusername": command{0, []string{"@username"}, "update @username", nil},
+	"setname":     command{1, []string{"first", "last"}, "update name", nil},
+	"setbio":      command{0, []string{"Lorem ipsum"}, "update about", nil},
+	"setpassword": command{0, []string{"old", "new"}, "set or remove password", nil},
+	"config":      command{0, []string{"param", "value"}, "view or update configuration options", nil},
+	"report":      command{2, []string{"chat", "comment"}, "report a chat by id or @username", nil},
+	"add":         command{1, []string{"@username"}, "add @username to your chat list", nil},
+	"join":        command{1, []string{"https://t.me/invite_link"}, "join to chat via invite link or @publicname", nil},
+	"supergroup":  command{1, []string{"title", "description"}, "create new supergroup «title» with «description»", nil},
+	"channel":     command{1, []string{"title", "description"}, "create new channel «title» with «description»", nil},
 }
 
+var notForGroups = []ChatType{ChatTypeBasicGroup, ChatTypeSupergroup, ChatTypeChannel}
+var notForPM = []ChatType{ChatTypePrivate, ChatTypeSecret}
+var onlyForSecret = []ChatType{ChatTypePrivate, ChatTypeBasicGroup, ChatTypeSupergroup, ChatTypeChannel}
+
 var chatCommands = map[string]command{
-	"help":       command{0, []string{}, "help"},
-	"d":          command{0, []string{"n"}, "delete your last message(s)"},
-	"s":          command{1, []string{"edited message"}, "edit your last message"},
-	"silent":     command{1, []string{"message"}, "send a message without sound"},
-	"schedule":   command{2, []string{"{online | 2006-01-02T15:04:05 | 15:04:05}", "message"}, "schedules a message either to timestamp or to whenever the user goes online"},
-	"forward":    command{2, []string{"message_id", "target_chat"}, "forwards a message"},
-	"vcard":      command{0, []string{}, "print vCard as text"},
-	"add":        command{1, []string{"@username"}, "add @username to your chat list"},
-	"join":       command{1, []string{"https://t.me/invite_link"}, "join to chat via invite link or @publicname"},
-	"group":      command{1, []string{"title"}, "create groupchat «title» with current user"},
-	"supergroup": command{1, []string{"title", "description"}, "create new supergroup «title» with «description»"},
-	"channel":    command{1, []string{"title", "description"}, "create new channel «title» with «description»"},
-	"secret":     command{0, []string{}, "create secretchat with current user"},
-	"search":     command{0, []string{"string", "[limit]"}, "search <string> in current chat"},
-	"history":    command{0, []string{"limit"}, "get last [limit] messages from current chat"},
-	"block":      command{0, []string{}, "blacklist current user"},
-	"unblock":    command{0, []string{}, "unblacklist current user"},
-	"invite":     command{1, []string{"id or @username"}, "add user to current chat"},
-	"link":       command{0, []string{}, "get invite link for current chat"},
-	"kick":       command{1, []string{"id or @username"}, "remove user to current chat"},
-	"mute":       command{1, []string{"id or @username", "hours"}, "mute user in current chat"},
-	"unmute":     command{1, []string{"id or @username"}, "unrestrict user from current chat"},
-	"ban":        command{1, []string{"id or @username", "hours"}, "restrict @username from current chat for [hours] or forever"},
-	"unban":      command{1, []string{"id or @username"}, "unbans @username in current chat (and devotes from admins)"},
-	"promote":    command{1, []string{"id or @username", "title"}, "promote user to admin in current chat"},
-	"leave":      command{0, []string{}, "leave current chat"},
-	"leave!":     command{0, []string{}, "leave current chat (for owners)"},
-	"ttl":        command{0, []string{"seconds"}, "set secret chat messages TTL before self-destroying"},
-	"close":      command{0, []string{}, "close current secret chat"},
-	"delete":     command{0, []string{}, "delete current chat from chat list"},
-	"members":    command{0, []string{"query"}, "search members [by optional query] in current chat (requires admin rights)"},
+	"help":       command{0, []string{}, "help", nil},
+	"d":          command{0, []string{"n"}, "delete your last message(s)", nil},
+	"s":          command{1, []string{"edited message"}, "edit your last message", nil},
+	"silent":     command{1, []string{"message"}, "send a message without sound", nil},
+	"schedule":   command{2, []string{"{online | 2006-01-02T15:04:05 | 15:04:05}", "message"}, "schedules a message either to timestamp or to whenever the user goes online", nil},
+	"forward":    command{2, []string{"message_id", "target_chat"}, "forwards a message", nil},
+	"vcard":      command{0, []string{}, "print vCard as text", nil},
+	"add":        command{1, []string{"@username"}, "add @username to your chat list", nil},
+	"join":       command{1, []string{"https://t.me/invite_link"}, "join to chat via invite link or @publicname", nil},
+	"group":      command{1, []string{"title"}, "create groupchat «title» with current user", &notForGroups},
+	"supergroup": command{1, []string{"title", "description"}, "create new supergroup «title» with «description»", nil},
+	"channel":    command{1, []string{"title", "description"}, "create new channel «title» with «description»", nil},
+	"secret":     command{0, []string{}, "create secretchat with current user", &notForGroups},
+	"search":     command{0, []string{"string", "[limit]"}, "search <string> in current chat", nil},
+	"history":    command{0, []string{"limit"}, "get last [limit] messages from current chat", nil},
+	"block":      command{0, []string{}, "blacklist current user", &notForGroups},
+	"unblock":    command{0, []string{}, "unblacklist current user", &notForGroups},
+	"invite":     command{1, []string{"id or @username"}, "add user to current chat", &notForPM},
+	"link":       command{0, []string{}, "get invite link for current chat", &notForPM},
+	"kick":       command{1, []string{"id or @username"}, "remove user from current chat", &notForPM},
+	"mute":       command{1, []string{"id or @username", "hours"}, "mute user in current chat", &notForPM},
+	"unmute":     command{1, []string{"id or @username"}, "unrestrict user from current chat", &notForPM},
+	"ban":        command{1, []string{"id or @username", "hours"}, "restrict @username from current chat for [hours] or forever", &notForPM},
+	"unban":      command{1, []string{"id or @username"}, "unbans @username in current chat (and devotes from admins)", &notForPM},
+	"promote":    command{1, []string{"id or @username", "title"}, "promote user to admin in current chat", &notForPM},
+	"leave":      command{0, []string{}, "leave current chat", &notForPM},
+	"leave!":     command{0, []string{}, "leave current chat (for owners)", &notForPM},
+	"ttl":        command{0, []string{"seconds"}, "set secret chat messages TTL before self-destroying", &onlyForSecret},
+	"close":      command{0, []string{}, "close current secret chat", &onlyForSecret},
+	"delete":     command{0, []string{}, "delete current chat from chat list", nil},
+	"members":    command{0, []string{"query"}, "search members [by optional query] in current chat (requires admin rights)", nil},
 }
 
 var transportConfigurationOptions = map[string]configurationOption{
@@ -112,6 +116,7 @@ type command struct {
 	RequiredArgs int
 	Arguments    []string
 	Description  string
+	NotFor       *[]ChatType
 }
 type configurationOption struct {
 	arguments   string
@@ -185,14 +190,31 @@ func CommandToHelpString(name string, cmd command) string {
 	return str.String()
 }
 
-func helpString(typ CommandType) string {
+// IsCommandFor checks the suitability of a command for a chat type
+func IsCommandForChatType(cmd command, chatType ChatType) bool {
+	if cmd.NotFor != nil {
+		for _, typ := range *cmd.NotFor {
+			if chatType == typ {
+				return false
+			}
+		}
+	}
+
+	return true
+}
+
+func (c *Client) helpString(typ CommandType, chatId int64) string {
 	var str strings.Builder
 
 	commandMap := GetCommands(typ)
+	chatType, chatTypeErr := c.GetChatType(chatId)
 
 	str.WriteString("Available commands:\n")
 	for _, name := range SortedCommandKeys(commandMap) {
 		command := commandMap[name]
+		if chatTypeErr == nil && !IsCommandForChatType(command, chatType) {
+			continue
+		}
 		str.WriteString(CommandToHelpString(name, command))
 		str.WriteString("\n")
 	}
@@ -502,7 +524,7 @@ func (c *Client) ProcessTransportCommand(cmdline string, resource string) string
 	case "channel":
 		return c.cmdChannel(args, cmdline)
 	case "help":
-		return helpString(CommandTypeTransport)
+		return c.helpString(CommandTypeTransport, 0)
 	}
 
 	return ""
@@ -524,6 +546,11 @@ func (c *Client) ProcessChatCommand(chatID int64, cmdline string) (string, bool)
 		return notEnoughArguments, true
 	}
 
+	chatType, chatTypeErr := c.GetChatType(chatID)
+	if chatTypeErr == nil && !IsCommandForChatType(command, chatType) {
+		return "Not applicable for this chat type", true
+	}
+
 	switch cmd {
 	// delete message
 	case "d":
@@ -1103,7 +1130,7 @@ func (c *Client) ProcessChatCommand(chatID int64, cmdline string) (string, bool)
 
 		return strings.Join(entries, "\n"), true
 	case "help":
-		return helpString(CommandTypeChat), true
+		return c.helpString(CommandTypeChat, chatID), true
 	default:
 		return "", false
 	}
diff --git a/telegram/utils.go b/telegram/utils.go
index 7ab5765..2c8b00d 100644
--- a/telegram/utils.go
+++ b/telegram/utils.go
@@ -53,6 +53,18 @@ var replyRegex = regexp.MustCompile("\\A>>? ?([0-9]+)\\n")
 const newlineChar string = "\n"
 const messageHeaderSeparator string = " | " // no hrunicode allowed here yet
 
+// ChatType is an enum of chat types, roughly corresponding to TDLib's one but better
+type ChatType int
+
+const (
+	ChatTypeUnknown ChatType = iota
+	ChatTypePrivate
+	ChatTypeBasicGroup
+	ChatTypeSupergroup
+	ChatTypeSecret
+	ChatTypeChannel
+)
+
 // GetContactByUsername resolves username to user id retrieves user and chat information
 func (c *Client) GetContactByUsername(username string) (*client.Chat, *client.User, error) {
 	if !c.Online() {
@@ -130,10 +142,10 @@ func (c *Client) GetContactByID(id int64, chat *client.Chat) (*client.Chat, *cli
 	return chat, user, nil
 }
 
-// IsPM checks if a chat is PM
-func (c *Client) IsPM(id int64) (bool, error) {
+// GetChatType obtains chat type from its information
+func (c *Client) GetChatType(id int64) (ChatType, error) {
 	if !c.Online() || id == 0 {
-		return false, errOffline
+		return ChatTypeUnknown, errOffline
 	}
 
 	var err error
@@ -144,14 +156,38 @@ func (c *Client) IsPM(id int64) (bool, error) {
 			ChatId: id,
 		})
 		if err != nil {
-			return false, err
+			return ChatTypeUnknown, err
 		}
 
 		c.cache.SetChat(id, chat)
 	}
 
 	chatType := chat.Type.ChatTypeType()
-	if chatType == client.TypeChatTypePrivate || chatType == client.TypeChatTypeSecret {
+	if chatType == client.TypeChatTypePrivate {
+		return ChatTypePrivate, nil
+	} else if chatType == client.TypeChatTypeBasicGroup {
+		return ChatTypeBasicGroup, nil
+	} else if chatType == client.TypeChatTypeSupergroup {
+		supergroup, _ := chat.Type.(*client.ChatTypeSupergroup)
+		if supergroup.IsChannel {
+			return ChatTypeChannel, nil
+		}
+		return ChatTypeSupergroup, nil
+	} else if chatType == client.TypeChatTypeSecret {
+		return ChatTypeSecret, nil
+	}
+
+	return ChatTypeUnknown, errors.New("Unknown chat type")
+}
+
+// IsPM checks if a chat is PM
+func (c *Client) IsPM(id int64) (bool, error) {
+	typ, err := c.GetChatType(id)
+	if err != nil {
+		return false, err
+	}
+
+	if typ == ChatTypePrivate || typ == ChatTypeSecret {
 		return true, nil
 	}
 	return false, nil
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index c50dd1c..3554394 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -475,6 +475,21 @@ func handleGetVcardIq(s xmpp.Sender, iq *stanza.IQ, typ byte) {
 	_ = gateway.ResumableSend(component, &answer)
 }
 
+func getTelegramChatType(from string, to string) (telegram.ChatType, error) {
+	toId, ok := toToID(to)
+	if ok {
+		bare, _, ok := gateway.SplitJID(from)
+		if ok {
+			session, ok := sessions[bare]
+			if ok {
+				return session.GetChatType(toId)
+			}
+		}
+	}
+
+	return telegram.ChatTypeUnknown, errors.New("Unknown chat type")
+}
+
 func handleGetDiscoInfo(s xmpp.Sender, iq *stanza.IQ, di *stanza.DiscoInfo) {
 	answer, err := stanza.NewIQ(stanza.Attrs{
 		Type: stanza.IQTypeResult,
@@ -501,6 +516,8 @@ func handleGetDiscoInfo(s xmpp.Sender, iq *stanza.IQ, di *stanza.DiscoInfo) {
 		}
 		disco.AddFeatures(NSCommand)
 	} else {
+		chatType, chatTypeErr := getTelegramChatType(iq.From, iq.To)
+
 		var cmdType telegram.CommandType
 		if ok {
 			cmdType = telegram.CommandTypeChat
@@ -510,6 +527,9 @@ func handleGetDiscoInfo(s xmpp.Sender, iq *stanza.IQ, di *stanza.DiscoInfo) {
 
 		for name, command := range telegram.GetCommands(cmdType) {
 			if di.Node == name {
+				if chatTypeErr == nil && !telegram.IsCommandForChatType(command, chatType) {
+					break
+				}
 				answer.Payload = di
 				di.AddIdentity(telegram.CommandToHelpString(name, command), "automation", "command-node")
 				di.AddFeatures(NSCommand, "jabber:x:data")
@@ -549,6 +569,8 @@ func handleGetDiscoItems(s xmpp.Sender, iq *stanza.IQ, di *stanza.DiscoItems) {
 	if di.Node == NSCommand {
 		answer.Payload = di
 
+		chatType, chatTypeErr := getTelegramChatType(iq.From, iq.To)
+
 		var cmdType telegram.CommandType
 		if ok {
 			cmdType = telegram.CommandTypeChat
@@ -559,6 +581,9 @@ func handleGetDiscoItems(s xmpp.Sender, iq *stanza.IQ, di *stanza.DiscoItems) {
 		commands := telegram.GetCommands(cmdType)
 		for _, name := range telegram.SortedCommandKeys(commands) {
 			command := commands[name]
+			if chatTypeErr == nil && !telegram.IsCommandForChatType(command, chatType) {
+				continue
+			}
 			di.AddItem(iq.To, name, telegram.CommandToHelpString(name, command))
 		}
 	} else {