commands.go 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285
  1. // mautrix-discord - A Matrix-Discord puppeting bridge.
  2. // Copyright (C) 2022 Tulir Asokan
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU Affero General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // This program is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU Affero General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU Affero General Public License
  15. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  16. package main
  17. import (
  18. "context"
  19. "fmt"
  20. "strings"
  21. "github.com/skip2/go-qrcode"
  22. "maunium.net/go/mautrix/bridge/commands"
  23. "maunium.net/go/mautrix/event"
  24. "maunium.net/go/mautrix/id"
  25. "go.mau.fi/mautrix-discord/remoteauth"
  26. )
  27. type WrappedCommandEvent struct {
  28. *commands.Event
  29. Bridge *DiscordBridge
  30. User *User
  31. Portal *Portal
  32. }
  33. func (br *DiscordBridge) RegisterCommands() {
  34. proc := br.CommandProcessor.(*commands.Processor)
  35. proc.AddHandlers(
  36. cmdLogin,
  37. cmdLogout,
  38. cmdReconnect,
  39. cmdDisconnect,
  40. cmdGuilds,
  41. )
  42. }
  43. func wrapCommand(handler func(*WrappedCommandEvent)) func(*commands.Event) {
  44. return func(ce *commands.Event) {
  45. user := ce.User.(*User)
  46. var portal *Portal
  47. if ce.Portal != nil {
  48. portal = ce.Portal.(*Portal)
  49. }
  50. br := ce.Bridge.Child.(*DiscordBridge)
  51. handler(&WrappedCommandEvent{ce, br, user, portal})
  52. }
  53. }
  54. var cmdLogin = &commands.FullHandler{
  55. Func: wrapCommand(fnLogin),
  56. Name: "login",
  57. Help: commands.HelpMeta{
  58. Section: commands.HelpSectionAuth,
  59. Description: "Link the bridge to your Discord account by scanning a QR code.",
  60. },
  61. }
  62. func fnLogin(ce *WrappedCommandEvent) {
  63. if ce.User.IsLoggedIn() {
  64. ce.Reply("You're already logged in")
  65. return
  66. }
  67. client, err := remoteauth.New()
  68. if err != nil {
  69. ce.Reply("Failed to prepare login: %v", err)
  70. return
  71. }
  72. qrChan := make(chan string)
  73. doneChan := make(chan struct{})
  74. var qrCodeEvent id.EventID
  75. go func() {
  76. code := <-qrChan
  77. resp := sendQRCode(ce, code)
  78. qrCodeEvent = resp
  79. }()
  80. ctx := context.Background()
  81. if err = client.Dial(ctx, qrChan, doneChan); err != nil {
  82. close(qrChan)
  83. close(doneChan)
  84. ce.Reply("Error connecting to login websocket: %v", err)
  85. return
  86. }
  87. <-doneChan
  88. if qrCodeEvent != "" {
  89. _, _ = ce.MainIntent().RedactEvent(ce.RoomID, qrCodeEvent)
  90. }
  91. user, err := client.Result()
  92. if err != nil || len(user.Token) == 0 {
  93. ce.Reply("Error logging in: %v", err)
  94. } else if err = ce.User.Login(user.Token); err != nil {
  95. ce.Reply("Error connecting after login: %v", err)
  96. }
  97. ce.User.Lock()
  98. ce.User.ID = user.UserID
  99. ce.User.Update()
  100. ce.User.Unlock()
  101. ce.Reply("Successfully logged in as %s#%s", user.Username, user.Discriminator)
  102. }
  103. func sendQRCode(ce *WrappedCommandEvent, code string) id.EventID {
  104. url, ok := uploadQRCode(ce, code)
  105. if !ok {
  106. return ""
  107. }
  108. content := event.MessageEventContent{
  109. MsgType: event.MsgImage,
  110. Body: code,
  111. URL: url.CUString(),
  112. }
  113. resp, err := ce.Bot.SendMessageEvent(ce.RoomID, event.EventMessage, &content)
  114. if err != nil {
  115. ce.Log.Errorfln("Failed to send QR code: %v", err)
  116. return ""
  117. }
  118. return resp.EventID
  119. }
  120. func uploadQRCode(ce *WrappedCommandEvent, code string) (id.ContentURI, bool) {
  121. qrCode, err := qrcode.Encode(code, qrcode.Low, 256)
  122. if err != nil {
  123. ce.Log.Errorln("Failed to encode QR code:", err)
  124. ce.Reply("Failed to encode QR code: %v", err)
  125. return id.ContentURI{}, false
  126. }
  127. resp, err := ce.Bot.UploadBytes(qrCode, "image/png")
  128. if err != nil {
  129. ce.Log.Errorln("Failed to upload QR code:", err)
  130. ce.Reply("Failed to upload QR code: %v", err)
  131. return id.ContentURI{}, false
  132. }
  133. return resp.ContentURI, true
  134. }
  135. var cmdLogout = &commands.FullHandler{
  136. Func: wrapCommand(fnLogout),
  137. Name: "logout",
  138. Help: commands.HelpMeta{
  139. Section: commands.HelpSectionAuth,
  140. Description: "Unlink the bridge from your WhatsApp account.",
  141. },
  142. RequiresLogin: true,
  143. }
  144. func fnLogout(ce *WrappedCommandEvent) {
  145. err := ce.User.Logout()
  146. if err != nil {
  147. ce.Reply("Error logging out: %v", err)
  148. } else {
  149. ce.Reply("Logged out successfully.")
  150. }
  151. }
  152. var cmdDisconnect = &commands.FullHandler{
  153. Func: wrapCommand(fnDisconnect),
  154. Name: "disconnect",
  155. Help: commands.HelpMeta{
  156. Section: commands.HelpSectionAuth,
  157. Description: "Disconnect from Discord (without logging out)",
  158. },
  159. RequiresLogin: true,
  160. }
  161. func fnDisconnect(ce *WrappedCommandEvent) {
  162. if !ce.User.Connected() {
  163. ce.Reply("You're already not connected")
  164. } else if err := ce.User.Disconnect(); err != nil {
  165. ce.Reply("Error while disconnecting: %v", err)
  166. } else {
  167. ce.Reply("Successfully disconnected")
  168. }
  169. }
  170. var cmdReconnect = &commands.FullHandler{
  171. Func: wrapCommand(fnReconnect),
  172. Name: "reconnect",
  173. Aliases: []string{"connect"},
  174. Help: commands.HelpMeta{
  175. Section: commands.HelpSectionAuth,
  176. Description: "Reconnect to Discord after disconnecting",
  177. },
  178. RequiresLogin: true,
  179. }
  180. func fnReconnect(ce *WrappedCommandEvent) {
  181. if ce.User.Connected() {
  182. ce.Reply("You're already connected")
  183. } else if err := ce.User.Connect(); err != nil {
  184. ce.Reply("Error while reconnecting: %v", err)
  185. } else {
  186. ce.Reply("Successfully reconnected")
  187. }
  188. }
  189. var cmdGuilds = &commands.FullHandler{
  190. Func: wrapCommand(fnGuilds),
  191. Name: "guilds",
  192. Aliases: []string{"servers", "guild", "server"},
  193. Help: commands.HelpMeta{
  194. Section: commands.HelpSectionUnclassified,
  195. Description: "Guild bridging management",
  196. Args: "<status/bridge/unbridge> [_guild ID_] [--entire]",
  197. },
  198. RequiresLogin: true,
  199. }
  200. func fnGuilds(ce *WrappedCommandEvent) {
  201. if len(ce.Args) == 0 {
  202. ce.Reply("**Usage**: `$cmdprefix guilds <status/bridge/unbridge> [guild ID] [--entire]`")
  203. }
  204. subcommand := strings.ToLower(ce.Args[0])
  205. ce.Args = ce.Args[1:]
  206. switch subcommand {
  207. case "status":
  208. fnListGuilds(ce)
  209. case "bridge":
  210. fnBridgeGuild(ce)
  211. case "unbridge":
  212. fnUnbridgeGuild(ce)
  213. }
  214. }
  215. func fnListGuilds(ce *WrappedCommandEvent) {
  216. ce.User.guildsLock.Lock()
  217. defer ce.User.guildsLock.Unlock()
  218. if len(ce.User.guilds) == 0 {
  219. ce.Reply("You haven't joined any guilds")
  220. } else {
  221. var output strings.Builder
  222. for _, guild := range ce.User.guilds {
  223. status := "not bridged"
  224. if guild.Bridge {
  225. status = "bridged"
  226. }
  227. _, _ = fmt.Fprintf(&output, "* %s (`%s`) - %s\n", guild.GuildName, guild.GuildID, status)
  228. }
  229. ce.Reply("List of guilds:\n\n%s", output.String())
  230. }
  231. }
  232. func fnBridgeGuild(ce *WrappedCommandEvent) {
  233. if len(ce.Args) == 0 || len(ce.Args) > 2 {
  234. ce.Reply("**Usage**: `$cmdprefix guilds bridge <guild ID> [--entire]")
  235. } else if err := ce.User.bridgeGuild(ce.Args[0], len(ce.Args) == 2 && strings.ToLower(ce.Args[1]) == "--entire"); err != nil {
  236. ce.Reply("Error bridging guild: %v", err)
  237. } else {
  238. ce.Reply("Successfully bridged guild")
  239. }
  240. }
  241. func fnUnbridgeGuild(ce *WrappedCommandEvent) {
  242. if len(ce.Args) != 1 {
  243. ce.Reply("**Usage**: `$cmdprefix guilds unbridge <guild ID>")
  244. } else if err := ce.User.unbridgeGuild(ce.Args[0]); err != nil {
  245. ce.Reply("Error unbridging guild: %v", err)
  246. } else {
  247. ce.Reply("Successfully unbridged guild")
  248. }
  249. }