package main

import (
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/kataras/jwt"
)

// The jwt package provides an extra feature to protect your claims,
// you can declare what fields of your custom claims structure are required
// (navigate through custom-validations example too).
// Useful when the same key can generate more than one token claims type
// and you want somehow to separate the token claims type (e.g. a refresh token and a access claims).
//
// [1] Just change the default Unmarshal function with the following:
func init() {
	jwt.Unmarshal = jwt.UnmarshalWithRequired
}

// [2] And add a ',required' after your json tags:
type userClaims struct {
	Expiry   int64  `json:"exp"`
	Username string `json:"username,required"`
}

func main() {
	http.HandleFunc("/", getTokenHandler)
	http.HandleFunc("/protected", verifyTokenHandler)

	http.HandleFunc("/user", getUserTokenHandler)
	http.HandleFunc("/user/protected", verifyUserTokenHandler)

	log.Printf("Server listening on: http://localhost:8080")

	// http://localhost:8080
	// http://localhost:8080/protected?token={token} (OK)
	// http://localhost:8080/user
	// http://localhost:8080/user/protected?token={user_token} (OK)
	// [3] http://localhost:8080/user/protected?token={token} (NOT OK)
	http.ListenAndServe(":8080", nil)
}

var sharedKey = []byte("sercrethatmaycontainch@r$32chars")

// generate token to use.
func getTokenHandler(w http.ResponseWriter, r *http.Request) {
	claims := jwt.Map{"foo": "bar"} // the "username" is missing.

	token, err := jwt.Sign(jwt.HS256, sharedKey, claims)
	if err != nil {
		log.Printf("Generate token failure: %v", err)
		http.Error(w, "failure: sign and encode the token", http.StatusInternalServerError)
		return
	}

	tokenString := jwt.BytesToString(token)

	w.Header().Set("Content-Type", "text/html;charset=utf-8")
	fmt.Fprintf(w, `Token: %s<br/><br/><a href="/protected?token=%s">/protected?token=%s</a><br/>`,
		tokenString, tokenString, tokenString)

	// Add second link of the User Token, it should be result on 401:
	fmt.Fprintf(w, `<br/> This shows the required usage, try to access the /user with the current "/" token:<br/><a href="/user/protected?token=%s">/user/protected?token=%s</a>`,
		tokenString, tokenString)
}

func verifyTokenHandler(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("token")
	if token == "" {
		log.Printf("Token is missing")
		unauthorized(w)
		return
	}

	verifiedToken, err := jwt.Verify(jwt.HS256, sharedKey, []byte(token))
	if err != nil {
		log.Printf("Verify error: %v", err)
		unauthorized(w)
		return
	}

	var claims jwt.Map
	if err = verifiedToken.Claims(&claims); err != nil {
		log.Printf("Verify: decode claims: %v", err)
		unauthorized(w)
		return
	}

	fmt.Fprintf(w, "This is an authenticated request made of token: %q\n\n", token)
	fmt.Fprintf(w, "Claims:\n%#+v\n", claims)
	fmt.Fprintf(w, "Standard Claims:\n%#+v", verifiedToken.StandardClaims)
}

// generate token to use.
func getUserTokenHandler(w http.ResponseWriter, r *http.Request) {
	claims := userClaims{
		Expiry:   time.Now().Add(1 * time.Minute).Unix(),
		Username: "kataras",
	}

	token, err := jwt.Sign(jwt.HS256, sharedKey, claims)
	if err != nil {
		log.Printf("Generate token failure: %v", err)
		http.Error(w, "failure: sign and encode the token", http.StatusInternalServerError)
		return
	}

	tokenString := jwt.BytesToString(token)

	w.Header().Set("Content-Type", "text/html;charset=utf-8")
	fmt.Fprintf(w, `Token: %s<br/><br/><a href="/user/protected?token=%s">/user/protected?token=%s</a>`,
		tokenString, tokenString, tokenString)
}

func verifyUserTokenHandler(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("token")
	if token == "" {
		log.Printf("Token is missing")
		unauthorized(w)
		return
	}

	verifiedToken, err := jwt.Verify(jwt.HS256, sharedKey, []byte(token))
	if err != nil {
		log.Printf("Verify error: %v", err)
		unauthorized(w)
		return
	}

	var claims userClaims
	if err = verifiedToken.Claims(&claims); err != nil {
		// If a different token, e.g. generated by root "/" path
		// is passed and a "username" field is missing then:
		// Verify: decode claims: token is missing a required field: "Username"
		log.Printf("Verify: decode claims: %v", err)
		unauthorized(w)
		return
	}

	fmt.Fprintf(w, "This is an authenticated request made of token: %q\n\n", token)
	fmt.Fprintf(w, "Claims:\n%#+v\n", claims)
	fmt.Fprintf(w, "Standard Claims:\n%#+v", verifiedToken.StandardClaims)
}

func unauthorized(w http.ResponseWriter) {
	http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
}