arimelody.me/api/v1/admin/admin.go

package admin

import (
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"arimelody.me/arimelody.me/discord"
)

type (
    Session struct {
        UserID string
        Token string
    }

    loginData struct {
        UserID string
        Password string
    }
)

const TOKEN_LENGTH = 64
const TOKEN_CHARS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

// TODO: consider relying *entirely* on env vars instead of hard-coded fallbacks
var ADMIN_ID_DISCORD = func() string {
    envvar := os.Getenv("DISCORD_ADMIN_ID")
    if envvar != "" {
        return envvar
    } else {
        return "356210742200107009"
    }
}()

var sessions []*Session

func CreateSession(UserID string) Session {
    return Session{
        UserID: UserID,
        Token: string(generateToken()),
    }
}

func Handler() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        fmt.Println(r.URL.Path)

        w.WriteHeader(200)
        w.Write([]byte("hello admin!"))
    })
}

func AuthorisedHandler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        auth := r.Header.Get("Authorization")
        if auth == "" || !strings.HasPrefix(auth, "Bearer ") {
            cookie, err := r.Cookie("token")
            if err != nil {
                w.WriteHeader(401)
                w.Write([]byte("Unauthorized"))
                return
            }
            auth = cookie.Value
        }
        auth = auth[7:]

        var session *Session
        for _, s := range sessions {
            if s.Token == auth {
                session = s
                break
            }
        } 

        if session == nil {
            w.WriteHeader(401)
            w.Write([]byte("Unauthorized"))
            return
        }

        ctx := context.WithValue(r.Context(), "token", session.Token)
        next.ServeHTTP(w, r.WithContext(ctx))
    })
}

func LoginHandler() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        code := r.URL.Query().Get("code")

        if code == "" {
            w.Header().Add("Location", discord.REDIRECT_URI)
            w.WriteHeader(307)
            return
        }

        // let's get an oauth token!
        req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/oauth2/token", discord.API_ENDPOINT),
            strings.NewReader(url.Values{
            "client_id": {discord.CLIENT_ID},
            "client_secret": {discord.CLIENT_SECRET},
            "grant_type": {"authorization_code"},
            "code": {code},
            "redirect_uri": {discord.MY_REDIRECT_URI},
        }.Encode()))
        req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

        res, err := http.DefaultClient.Do(req)
        if err != nil {
            fmt.Printf("Failed to retrieve OAuth token: %s\n", err)
            w.WriteHeader(500)
            w.Write([]byte("Internal server error"))
            return
        }

        oauth := discord.AccessTokenResponse{}

        err = json.NewDecoder(res.Body).Decode(&oauth)
        if err != nil {
            fmt.Printf("Failed to parse OAuth response data from discord: %s\n", err)
            w.WriteHeader(500)
            w.Write([]byte("Internal server error"))
            return
        }
        res.Body.Close()

        discord_access_token := oauth.AccessToken

        // let's get authorisation information!
        req, err = http.NewRequest(http.MethodGet, fmt.Sprintf("%s/oauth2/@me", discord.API_ENDPOINT), nil)
        req.Header.Add("Authorization", "Bearer " + discord_access_token)

        res, err = http.DefaultClient.Do(req)
        if err != nil {
            fmt.Printf("Failed to retrieve discord auth information: %s\n", err)
            w.WriteHeader(500)
            w.Write([]byte("Internal server error"))
            return
        }

        auth_info := discord.AuthInfoResponse{}

        err = json.NewDecoder(res.Body).Decode(&auth_info)
        if err != nil {
            fmt.Printf("Failed to parse auth information from discord: %s\n", err)
            w.WriteHeader(500)
            w.Write([]byte("Internal server error"))
            return
        }
        res.Body.Close()

        discord_user_id := auth_info.User.Id
        
        if discord_user_id != ADMIN_ID_DISCORD {
            // TODO: unauthorized user. revoke the token
            w.WriteHeader(401)
            w.Write([]byte("Unauthorized"))
            return
        }

        // login success!
        session := CreateSession(auth_info.User.Username)
        sessions = append(sessions, &session)

        cookie := http.Cookie{}
        cookie.Name = "token"
        cookie.Value = session.Token
        cookie.Expires = time.Now().Add(24 * time.Hour)
        // cookie.Secure = true
        cookie.HttpOnly = true
        cookie.Path = "/"
        http.SetCookie(w, &cookie)

        w.WriteHeader(200)
        w.Write([]byte(session.Token))
    })
}

func LogoutHandler() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        token := r.Context().Value("token").(string)

        if token == "" {
            w.WriteHeader(401)
            return
        }

        // remove this session from the list
        sessions = func (token string) []*Session {
            new_sessions := []*Session{}
            for _, session := range sessions {
                new_sessions = append(new_sessions, session)
            }
            return new_sessions
        }(token)

        w.WriteHeader(200)
    })
}

func OAuthCallbackHandler() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    })
}

func VerifyHandler() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // this is an authorised endpoint, so you *must* supply a valid token
        // before accessing this route.
        w.WriteHeader(200)
    })
}

func generateToken() string {
    var token []byte

    for i := 0; i < TOKEN_LENGTH; i++ {
        token = append(token, TOKEN_CHARS[rand.Intn(len(TOKEN_CHARS))])
    }

    return string(token)
}