package discord import ( "bytes" "code.mrmelon54.xyz/sean/melon-tools/utils" "context" "embed" _ "embed" "fmt" "github.com/bwmarrin/discordgo" "github.com/google/uuid" "github.com/gorilla/mux" "golang.org/x/oauth2" "html/template" "image" "image/png" "net/http" "os" ) var ( //go:embed pages/index.go.html indexTemplate string //go:embed pages/assets/icon iconFiles embed.FS ) type Module struct { sessionWrapper func(cb func(http.ResponseWriter, *http.Request, *utils.State)) func(rw http.ResponseWriter, req *http.Request) oauthClient *oauth2.Config } type discordKeyType int const ( KeyOauthClient = discordKeyType(iota) KeyUser KeyState KeyAccessToken KeyRefreshToken ) func New() *Module { return &Module{} } func (m *Module) GetName() string { return "Discord" } func (m *Module) GetEndpoint() string { return "/discord" } func (m *Module) SetupModule(router *mux.Router, f func(cb func(http.ResponseWriter, *http.Request, *utils.State)) func(rw http.ResponseWriter, req *http.Request)) { m.sessionWrapper = f m.oauthClient = &oauth2.Config{ ClientID: os.Getenv("DISCORD_CLIENT_ID"), ClientSecret: os.Getenv("DISCORD_CLIENT_SECRET"), Scopes: []string{"identify", "guilds", "connections", "email"}, Endpoint: oauth2.Endpoint{ AuthURL: "https://discord.com/oauth2/authorize", TokenURL: "https://discord.com/api/oauth2/token", }, RedirectURL: os.Getenv("DISCORD_REDIRECT_URL"), } router.HandleFunc("/", m.getClient(m.homepage)) router.HandleFunc("/login", m.sessionWrapper(m.loginPage)) router.HandleFunc("/user/avatar/{userId}/{avatarId}", m.getClient(m.userAvatar)) router.HandleFunc("/guild/icon/{guildId}/{iconId}", m.getClient(m.guildIcon)) router.PathPrefix("/assets/icon/{name}.svg").HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { vars := mux.Vars(req) b, err := iconFiles.ReadFile("pages/assets/icon/" + vars["name"] + ".svg") if err != nil { rw.WriteHeader(http.StatusNotFound) } else { rw.Header().Set("Content-Type", "image/svg+xml") rw.WriteHeader(http.StatusOK) _, _ = rw.Write(b) } }) } func (m *Module) getClient(cb func(http.ResponseWriter, *http.Request, *utils.State, *discordgo.Session)) func(rw http.ResponseWriter, req *http.Request) { return m.sessionWrapper(func(rw http.ResponseWriter, req *http.Request, state *utils.State) { if v, ok := utils.GetStateValue[*discordgo.Session](state, KeyOauthClient); ok { cb(rw, req, state, v) return } http.Redirect(rw, req, "/discord/login", http.StatusTemporaryRedirect) }) } func (m *Module) homepage(rw http.ResponseWriter, req *http.Request, state *utils.State, discordClient *discordgo.Session) { myUser, err := discordClient.User("@me") if err != nil { state.Del(KeyOauthClient) http.Error(rw, err.Error(), http.StatusInternalServerError) return } myGuilds, err := discordClient.UserGuilds(100, "", "") if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) return } myConns, err := discordClient.UserConnections() if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) return } tmp, err := template.New("homepage").Funcs(template.FuncMap{ "checkFlag": func(a discordgo.UserFlags, b int) bool { return (int(a) & b) != 0 }, "connectedLink": connectedLinkFunc, }).Parse(indexTemplate) if err != nil { fmt.Println("Template parse error:", err) return } guildIcons := make([]template.HTMLAttr, len(myGuilds)) for i, j := range myGuilds { guildIcons[i] = template.HTMLAttr(fmt.Sprintf("src=\"/discord/guild/icon/%s/%s\"", j.ID, j.Icon)) } err = tmp.Execute(rw, struct { User *discordgo.User UserAccent string Avatar template.HTMLAttr Guilds []*discordgo.UserGuild GuildIcons []template.HTMLAttr Connections []*discordgo.UserConnection }{ User: myUser, UserAccent: fmt.Sprintf("#%06x", myUser.AccentColor), Avatar: template.HTMLAttr(fmt.Sprintf("src=\"/discord/user/avatar/%s/%s\"", myUser.ID, myUser.Avatar)), Guilds: myGuilds, GuildIcons: guildIcons, Connections: myConns, }) if err != nil { fmt.Println("Template execute error:", err) return } } func (m *Module) loginPage(rw http.ResponseWriter, req *http.Request, state *utils.State) { if myUser, ok := utils.GetStateValue[*string](state, KeyUser); ok { if myUser != nil { http.Redirect(rw, req, "/discord", http.StatusTemporaryRedirect) return } } if flowState, ok := utils.GetStateValue[uuid.UUID](state, KeyState); ok { q := req.URL.Query() if q.Has("code") && q.Has("state") { if q.Get("state") == flowState.String() { exchange, err := m.oauthClient.Exchange(context.Background(), q.Get("code")) if err != nil { fmt.Println("Exchange token error:", err) return } c, err := discordgo.New("Bearer " + exchange.AccessToken) if err != nil { fmt.Println("Create client error:", err) return } state.Put(KeyOauthClient, c) state.Put(KeyAccessToken, exchange.AccessToken) state.Put(KeyRefreshToken, exchange.RefreshToken) http.Redirect(rw, req, "/discord", http.StatusTemporaryRedirect) return } http.Error(rw, "OAuth flow state doesn't match\n", http.StatusBadRequest) return } } flowState := uuid.New() state.Put(KeyState, flowState) http.Redirect(rw, req, m.oauthClient.AuthCodeURL(flowState.String(), oauth2.AccessTypeOffline), http.StatusTemporaryRedirect) } func (m *Module) userAvatar(rw http.ResponseWriter, req *http.Request, state *utils.State, discordClient *discordgo.Session) { vars := mux.Vars(req) body, err := discordClient.RequestWithBucketID("GET", discordgo.EndpointUserAvatar(vars["userId"], vars["avatarId"]), nil, discordgo.EndpointUserAvatar("", "")) if err != nil { return } rw.Header().Set("Content-Type", "image/png") rw.WriteHeader(200) myAvatar, _, err := image.Decode(bytes.NewReader(body)) err = png.Encode(rw, myAvatar) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) return } } func (m *Module) guildIcon(rw http.ResponseWriter, req *http.Request, state *utils.State, discordClient *discordgo.Session) { vars := mux.Vars(req) body, err := discordClient.RequestWithBucketID("GET", discordgo.EndpointGuildIcon(vars["guildId"], vars["iconId"]), nil, discordgo.EndpointGuildIcon(vars["guildId"], "")) if err != nil { return } rw.Header().Set("Content-Type", "image/png") rw.WriteHeader(http.StatusOK) myAvatar, _, err := image.Decode(bytes.NewReader(body)) err = png.Encode(rw, myAvatar) if err != nil { http.Error(rw, err.Error(), http.StatusInternalServerError) return } } func connectedLinkFunc(a *discordgo.UserConnection) string { switch a.Type { case "github": return "https://github.com/" + a.Name case "reddit": return "https://www.reddit.com/u/" + a.Name case "spotify": return "https://open.spotify.com/user/" + a.ID case "steam": return "https://steamcommunity.com/profiles/" + a.ID case "twitch": return "https://www.twitch.tv/" + a.Name case "twitter": return "https://twitter.com/" + a.Name case "xbox": return "javascript:alert('No link to XBox profiles')" case "youtube": return "https://www.youtube.com/channel/" + a.ID } return "javascript:alert('Unknown profile type')" }