package jwt

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"

	vault "github.com/hashicorp/vault/api"
)

// DefaultMountPath specifies the default mount path for the JWT
// Authentication Method.
const DefaultMountPath = "jwt"

// ErrNoToken is an error, which is returned when [JWTAuth] is configured
// with an empty token.
var ErrNoToken = errors.New("no token specified")

// ErrInvalidMountPath is an error, which is returned when configuring [JWTAuth]
// to use an invalid mount path for a Vault Authentication Method.
var ErrInvalidMountPath = errors.New("invalid auth method mount path specified")

// JWTAuth implements support for the [JWT Authentication Method].
//
// [JWT Authentication Method]: https://developer.hashicorp.com/vault/docs/auth/jwt
type JWTAuth struct {
	// roleName specifies the name of the role to use.
	roleName string

	// mountPath specifies the mount path for the JWT Authentication Method.
	mountPath string

	// token specifies the JWT token which will be used for authenticating
	// against the Vault Authentication Method endpoint.
	token string
}

var _ vault.AuthMethod = &JWTAuth{}

// Option is a function which configures [JWTAuth].
type Option func(a *JWTAuth) error

// New creates a new [JWTAuth] and configures it with the given options.
//
// The default mount path for the JWT Authentication Method is
// [DefaultMountPath]. In order to configure a different mount path for the
// Authentication Method you can use the [WithMountPath] option.
//
// The JWT token which will be used for authentication against the Vault
// Authentication Method login endpoint may be specified either as a string,
// from path, or via an environment variable. In order to configure the token
// for authentication use the [WithToken], [WithTokenFromPath] or
// [WithTokenFromEnv] options.
func New(roleName string, opts ...Option) (*JWTAuth, error) {
	jwtAuth := &JWTAuth{
		roleName:  roleName,
		mountPath: DefaultMountPath,
	}

	for _, opt := range opts {
		if err := opt(jwtAuth); err != nil {
			return nil, err
		}
	}

	if jwtAuth.token == "" {
		return nil, ErrNoToken
	}

	if jwtAuth.mountPath == "" {
		return nil, ErrInvalidMountPath
	}

	return jwtAuth, nil
}

// Login implements the [vault.AuthMethod] interface.
func (a *JWTAuth) Login(ctx context.Context, client *vault.Client) (*vault.Secret, error) {
	path := fmt.Sprintf("auth/%s/login", a.mountPath)
	data := map[string]any{
		"jwt":  a.token,
		"role": a.roleName,
	}

	return client.Logical().WriteWithContext(ctx, path, data)
}

// WithToken is an [Option], which configures [JWTAuth] to use the given token
// when authenticating against the Vault JWT Authentication Method.
func WithToken(token string) Option {
	opt := func(a *JWTAuth) error {
		a.token = token

		return nil
	}

	return opt
}

// WithTokenFromPath is an [Option], which configures [JWTAuth] to read the
// token from the given path.
func WithTokenFromPath(path string) Option {
	opt := func(a *JWTAuth) error {
		token, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		a.token = string(token)

		return nil
	}

	return opt
}

// WithTokenFromEnv is an [Option], which configures [JWTAuth] to read the token
// from the given environment variable.
func WithTokenFromEnv(env string) Option {
	opt := func(a *JWTAuth) error {
		value := os.Getenv(env)
		a.token = value

		return nil
	}

	return opt
}

// WithMountPath is an [Option], which configures [JWTAuth] to use the given
// mount path for the Vault Authentication Method.
func WithMountPath(mountPath string) Option {
	opt := func(a *JWTAuth) error {
		// Remove any trailing slashes from the given mount path
		a.mountPath = strings.TrimRight(mountPath, "/")

		return nil
	}

	return opt
}

package main

import (
	"context"
	"errors"
	"fmt"
	"os"

	vault "github.com/hashicorp/vault/api"

	jwtauth "vault-jwt-auth-method/pkg/auth/jwt"
)

func printErr(err error) {
	fmt.Println(err)
	os.Exit(1)
}

func main() {
	ctx := context.Background()

	// Create a JWT Authentication Method using the token from the following
	// role and token stored in the path
	authMethod, err := jwtauth.New(
		"my-role-name",
		jwtauth.WithTokenFromPath("/path/to/my/token"),
	)
	if err != nil {
		printErr(err)
	}

	// Create Vault client and authenticate using our JWT Authentication
	// Method implementation
	config := vault.DefaultConfig()
	client, err := vault.NewClient(config)
	if err != nil {
		printErr(err)
	}

	// Login using our Auth Method
	authSecret, err := client.Auth().Login(ctx, authMethod)
	if err != nil {
		printErr(err)
	}
	if authSecret == nil {
		printErr(errors.New("no auth info returned after login"))
	}

	// Read a sample KV v2 secret
	kvv2 := client.KVv2("kvv2")
	secret, err := kvv2.Get(ctx, "my/test/secret")
	if err != nil {
		printErr(err)
	}

	fmt.Println(secret.Data)
}