arimelody.me/admin/accounthttp.go

package admin

import (
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"arimelody-web/controller"
	"arimelody-web/model"

	"golang.org/x/crypto/bcrypt"
)

type TemplateData struct {
    Account *model.Account
    Message string
    Token string
}

func AccountHandler(app *model.AppState) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        account := r.Context().Value("account").(*model.Account)

        totps, err := controller.GetTOTPsForAccount(app.DB, account.ID)
        if err != nil {
            fmt.Printf("WARN: Failed to fetch TOTPs: %v\n", err)
            http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
        }

        type AccountResponse struct {
            Account *model.Account
            TOTPs []model.TOTP
        }

        err = pages["account"].Execute(w, AccountResponse{
            Account: account,
            TOTPs: totps,
        })
        if err != nil {
            fmt.Printf("WARN: Failed to render admin account page: %v\n", err)
            http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
        }
    })
}

func LoginHandler(app *model.AppState) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if r.Method == http.MethodGet {
            account, err := controller.GetAccountByRequest(app.DB, r)
            if err != nil {
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
                fmt.Fprintf(os.Stderr, "WARN: Failed to fetch account: %v\n", err)
                return
            }
            if account != nil {
                http.Redirect(w, r, "/admin", http.StatusFound)
                return
            }

            err = pages["login"].Execute(w, TemplateData{})
            if err != nil {
                fmt.Fprintf(os.Stderr, "WARN: Error rendering admin login page: %s\n", err)
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
                return
            }
            return
        }

        type LoginResponse struct {
            Account *model.Account
            Token string
            Message string
        }

        render := func(data LoginResponse) {
            err := pages["login"].Execute(w, data)
            if err != nil {
                fmt.Fprintf(os.Stderr, "WARN: Error rendering admin login page: %s\n", err)
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
                return
            }
        }

        if r.Method != http.MethodPost {
            http.NotFound(w, r);
            return
        }

        err := r.ParseForm()
        if err != nil {
            render(LoginResponse{ Message: "Malformed request." })
            return
        }

        type LoginRequest struct {
            Username string `json:"username"`
            Password string `json:"password"`
            TOTP string `json:"totp"`
        }
        credentials := LoginRequest{
            Username: r.Form.Get("username"),
            Password: r.Form.Get("password"),
            TOTP: r.Form.Get("totp"),
        }

        account, err := controller.GetAccount(app.DB, credentials.Username)
        if err != nil {
            render(LoginResponse{ Message: "Invalid username or password" })
            return
        }
        if account == nil {
            render(LoginResponse{ Message: "Invalid username or password" })
            return
        }

        err = bcrypt.CompareHashAndPassword([]byte(account.Password), []byte(credentials.Password))
        if err != nil {
            render(LoginResponse{ Message: "Invalid username or password" })
            return
        }

        totps, err := controller.GetTOTPsForAccount(app.DB, account.ID)
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to fetch TOTPs: %v\n", err)
            render(LoginResponse{ Message: "Something went wrong. Please try again." })
            return
        }
        if len(totps) > 0 {
            success := false
            for _, totp := range totps {
                check := controller.GenerateTOTP(totp.Secret, 0)
                if check == credentials.TOTP {
                    success = true
                    break
                }
            }
            if !success {
                render(LoginResponse{ Message: "Invalid TOTP" })
                return
            }
        } else {
            // TODO: user should be prompted to add 2FA method
        }

        // login success!
        token, err := controller.CreateToken(app.DB, account.ID, r.UserAgent())
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to create token: %v\n", err)
            render(LoginResponse{ Message: "Something went wrong. Please try again." })
            return
        }

        cookie := http.Cookie{}
        cookie.Name = model.COOKIE_TOKEN
        cookie.Value = token.Token
        cookie.Expires = token.ExpiresAt
        if strings.HasPrefix(app.Config.BaseUrl, "https") {
            cookie.Secure = true
        }
        cookie.HttpOnly = true
        cookie.Path = "/"
        http.SetCookie(w, &cookie)

        render(LoginResponse{ Account: account, Token: token.Token })
    })
}

func LogoutHandler(app *model.AppState) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if r.Method != http.MethodGet {
            http.NotFound(w, r)
            return
        }

        tokenStr := controller.GetTokenFromRequest(app.DB, r)

        if len(tokenStr) > 0 {
            err := controller.DeleteToken(app.DB, tokenStr)
            if err != nil {
                fmt.Fprintf(os.Stderr, "WARN: Failed to revoke token: %v\n", err)
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
                return
            }
        }

        cookie := http.Cookie{}
        cookie.Name = model.COOKIE_TOKEN
        cookie.Value = ""
        cookie.Expires = time.Now()
        if strings.HasPrefix(app.Config.BaseUrl, "https") {
            cookie.Secure = true
        }
        cookie.HttpOnly = true
        cookie.Path = "/"
        http.SetCookie(w, &cookie)
        http.Redirect(w, r, "/admin/login", http.StatusFound)
    })
}

func createAccountHandler(app *model.AppState) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        checkAccount, err := controller.GetAccountByRequest(app.DB, r)
        if err != nil {
            fmt.Printf("WARN: Failed to fetch account: %s\n", err)
            http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
            return
        }
        if checkAccount != nil {
            // user is already logged in
            http.Redirect(w, r, "/admin", http.StatusFound)
            return
        }

        type CreateAccountResponse struct {
            Account *model.Account
            Message string
        }

        render := func(data CreateAccountResponse) {
            err := pages["create-account"].Execute(w, data)
            if err != nil {
                fmt.Printf("WARN: Error rendering create account page: %s\n", err)
                http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
            }
        }

        if r.Method == http.MethodGet {
            render(CreateAccountResponse{})
            return
        }

        if r.Method != http.MethodPost {
            http.NotFound(w, r)
            return
        }

        err = r.ParseForm()
        if err != nil {
            render(CreateAccountResponse{
                Message: "Malformed data.",
            })
            return
        }

        type RegisterRequest struct {
            Username    string `json:"username"`
            Email       string `json:"email"`
            Password    string `json:"password"`
            Invite      string `json:"invite"`
        }
        credentials := RegisterRequest{
            Username: r.Form.Get("username"),
            Email: r.Form.Get("email"),
            Password: r.Form.Get("password"),
            Invite: r.Form.Get("invite"),
        }

        // make sure code exists in DB
        invite, err := controller.GetInvite(app.DB, credentials.Invite)
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to retrieve invite: %v\n", err)
            render(CreateAccountResponse{
                Message: "Something went wrong. Please try again.",
            })
            return
        }
        if invite == nil || time.Now().After(invite.ExpiresAt) {
            if invite != nil {
                err := controller.DeleteInvite(app.DB, invite.Code)
                if err != nil { fmt.Fprintf(os.Stderr, "WARN: Failed to delete expired invite: %v\n", err) }
            }
            render(CreateAccountResponse{
                Message: "Invalid invite code.",
            })
            return
        }

        hashedPassword, err := bcrypt.GenerateFromPassword([]byte(credentials.Password), bcrypt.DefaultCost)
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to generate password hash: %v\n", err)
            render(CreateAccountResponse{
                Message: "Something went wrong. Please try again.",
            })
            return
        }

        account := model.Account{
            Username: credentials.Username,
            Password: string(hashedPassword),
            Email: credentials.Email,
            AvatarURL: "/img/default-avatar.png",
        }
        err = controller.CreateAccount(app.DB, &account)
        if err != nil {
            if strings.HasPrefix(err.Error(), "pq: duplicate key") {
                render(CreateAccountResponse{
                    Message: "An account with that username already exists.",
                })
                return
            }
            fmt.Fprintf(os.Stderr, "WARN: Failed to create account: %v\n", err)
            render(CreateAccountResponse{
                Message: "Something went wrong. Please try again.",
            })
            return
        }

        err = controller.DeleteInvite(app.DB, invite.Code)
        if err != nil { fmt.Fprintf(os.Stderr, "WARN: Failed to delete expired invite: %v\n", err) }

        // registration success!
        token, err := controller.CreateToken(app.DB, account.ID, r.UserAgent())
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to create token: %v\n", err)
            // gracefully redirect user to login page
            http.Redirect(w, r, "/admin/login", http.StatusFound)
            return
        }

        cookie := http.Cookie{}
        cookie.Name = model.COOKIE_TOKEN
        cookie.Value = token.Token
        cookie.Expires = token.ExpiresAt
        if strings.HasPrefix(app.Config.BaseUrl, "https") {
            cookie.Secure = true
        }
        cookie.HttpOnly = true
        cookie.Path = "/"
        http.SetCookie(w, &cookie)

        err = pages["login"].Execute(w, TemplateData{
            Account: &account,
            Token: token.Token,
        })
        if err != nil {
            fmt.Fprintf(os.Stderr, "WARN: Failed to render login page: %v\n", err)
			http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
            return
        }
    })
}