package persistence

import (
	"crypto/aes"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/sessions"
	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/cookies"
	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/encryption"
)

// saveFunc performs a persistent store's save functionality using
// a key string, value []byte & (optional) expiration time.Duration
type saveFunc func(string, []byte, time.Duration) error

// loadFunc performs a load from a persistent store using a
// string key and returning the stored value as []byte
type loadFunc func(string) ([]byte, error)

// clearFunc performs a persistent store's clear functionality using
// a string key for the target of the deletion.
type clearFunc func(string) error

// initLockFunc returns a lock object for a persistent store using a
// string key
type initLockFunc func(string) sessions.Lock

// ticket is a structure representing the ticket used in server based
// session storage. It provides a unique per session decryption secret giving
// more security than the shared CookieSecret.
type ticket struct {
	id            string
	secret        []byte
	cookieBuilder cookies.Builder
}

// newTicket creates a new ticket. The ID & secret will be randomly created
// with 16 byte sizes. The ID will be prefixed & hex encoded.
func newTicket(cookieBuilder cookies.Builder) (*ticket, error) {
	rawID := make([]byte, 16)
	if _, err := io.ReadFull(rand.Reader, rawID); err != nil {
		return nil, fmt.Errorf("failed to create new ticket ID: %v", err)
	}
	// ticketID is hex encoded
	ticketID := fmt.Sprintf("%s-%s", cookieBuilder.GetName(), hex.EncodeToString(rawID))

	secret := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, secret); err != nil {
		return nil, fmt.Errorf("failed to create encryption secret: %v", err)
	}

	return &ticket{
		id:            ticketID,
		secret:        secret,
		cookieBuilder: cookieBuilder,
	}, nil
}

// encodeTicket encodes the Ticket to a string for usage in cookies
func (t *ticket) encodeTicket() string {
	return fmt.Sprintf("%s.%s", t.id, base64.RawURLEncoding.EncodeToString(t.secret))
}

// decodeTicket decodes an encoded ticket string
func decodeTicket(encTicket string, cookieBuilder cookies.Builder) (*ticket, error) {
	ticketParts := strings.Split(encTicket, ".")
	if len(ticketParts) != 2 {
		return nil, errors.New("failed to decode ticket")
	}
	ticketID, secretBase64 := ticketParts[0], ticketParts[1]

	secret, err := base64.RawURLEncoding.DecodeString(secretBase64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode encryption secret: %v", err)
	}

	return &ticket{
		id:            ticketID,
		secret:        secret,
		cookieBuilder: cookieBuilder,
	}, nil
}

// decodeTicketFromRequest retrieves a potential ticket cookie from a request
// and decodes it to a ticket.
func decodeTicketFromRequest(req *http.Request, cookieBuilder cookies.Builder) (*ticket, error) {
	val, err := cookieBuilder.ValidateRequest(req)
	if err != nil {
		return nil, fmt.Errorf("invalid cookie: %w", err)
	}

	// Valid cookie, decode the ticket
	return decodeTicket(val, cookieBuilder)
}

// saveSession encodes the SessionState with the ticket's secret and persists
// it to disk via the passed saveFunc.
func (t *ticket) saveSession(s *sessions.SessionState, saver saveFunc) error {
	c, err := t.makeCipher()
	if err != nil {
		return err
	}
	ciphertext, err := s.EncodeSessionState(c, false)
	if err != nil {
		return fmt.Errorf("failed to encode the session state with the ticket: %v", err)
	}
	return saver(t.id, ciphertext, t.cookieBuilder.GetExpiration())
}

// loadSession loads a session from the disk store via the passed loadFunc
// using the ticket.id as the key. It then decodes the SessionState using
// ticket.secret to make the AES-GCM cipher.
// finally it appends a lock implementation
func (t *ticket) loadSession(loader loadFunc, initLock initLockFunc) (*sessions.SessionState, error) {
	ciphertext, err := loader(t.id)
	if err != nil {
		return nil, fmt.Errorf("failed to load the session state with the ticket: %v", err)
	}
	c, err := t.makeCipher()
	if err != nil {
		return nil, err
	}

	sessionState, err := sessions.DecodeSessionState(ciphertext, c, false)
	if err != nil {
		return nil, err
	}
	lock := initLock(t.id)
	sessionState.Lock = lock
	return sessionState, nil
}

// clearSession uses the passed clearFunc to delete a session stored with a
// key of ticket.id
func (t *ticket) clearSession(clearer clearFunc) error {
	return clearer(t.id)
}

// setCookie sets the encoded ticket as a cookie
func (t *ticket) setCookie(rw http.ResponseWriter, req *http.Request, s *sessions.SessionState) error {
	ticketCookie, err := t.cookieBuilder.
		WithSignedValue(true).
		WithStart(*s.CreatedAt).
		MakeCookie(req, t.encodeTicket())
	if err != nil {
		return err
	}

	http.SetCookie(rw, ticketCookie)
	return nil
}

// clearCookie removes any cookies that would be where this ticket
// would set them
func (t *ticket) clearCookie(rw http.ResponseWriter, req *http.Request) error {
	cookie, err := t.cookieBuilder.
		WithExpiration(time.Hour*-1). // Set the expiration to negative to clear the cookie
		MakeCookie(req, "")
	if err != nil {
		return fmt.Errorf("could not clear cookie: %v", err)
	}

	http.SetCookie(rw, cookie)
	return nil
}

// makeCipher makes a AES-GCM cipher out of the ticket's secret
func (t *ticket) makeCipher() (encryption.Cipher, error) {
	c, err := encryption.NewGCMCipher(t.secret)
	if err != nil {
		return nil, fmt.Errorf("failed to make an AES-GCM cipher from the ticket secret: %v", err)
	}
	return c, nil
}