tools/module/discord/main.go

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')"
}