package models

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/bson"
)

// SessionKey represents a public and sometimes private key-pair for a user
// that will be stored on the user's record in the user store. These keys are
// used for signing authentication JWTs.
//
// This object is designed to be serialized to BSON. Other serializations can
// be added in the future as needed.
//
// There are two flavors of this record. A record with a private key (which
// implies a public key) is a key that the service generated and is used by the
// service to sign JWTs for the user. The private key is never given to the
// user. The private key is only used in the CreateToken flow, never the Verify
// flow. Currently (as of Nov 2021) the application sets a near-future NotAfter
// date and these get garbage collected. It might be nice to re-use them in the
// future for a while but it's not all that important.
//
// The other flavor of this key will have a public key but no private key.
// These are service keys. Service keys are given to programmatic actors that
// need to be able to mint their own JWTs for authentication to the service.
// For these keys the client will construct their own JWT and sign it with the
// private key and the service will validate the signature with the public key.
// These keys (as of Nov 2021) do not expire, though they can be revoked.
type SessionKey struct {
	KeyId       string
	Description string
	Revoked     *time.Time
	NotAfter    *time.Time
	NotBefore   *time.Time
	PublicKey   crypto.PublicKey
	PrivateKey  *ecdsa.PrivateKey
}

func GenerateSessionKey(ttl time.Duration) (*SessionKey, error) {
	pk, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	key := make([]byte, 8)
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}

	now := time.Now()
	notAfter := now.Add(ttl)

	return &SessionKey{
		KeyId:      hex.EncodeToString(key),
		Revoked:    nil,
		NotAfter:   &notAfter,
		NotBefore:  &now,
		PublicKey:  pk.Public(),
		PrivateKey: pk,
	}, nil
}

// IsGarbage checks to determine if a key is garbage that should be collected.
// The definition of garbage is similar to the inversion of the definition of
// vaild but revoked keys are not considered to be garbage since they may be
// useful for auditing later. Also keys that are not yet valid are not garbage.
func (s *SessionKey) IsGarbage() bool {
	if s.Revoked != nil {
		return false
	}

	if s.NotBefore != nil && time.Now().Before(*s.NotBefore) {
		return false
	}

	if s.NotAfter != nil && time.Now().After(*s.NotAfter) {
		return true
	}

	return false
}

// IsValid checks the various dates in the SessionKey to verify that they are
// valid and in-range for use. This should be called before trusting this key
// for any use.
func (s *SessionKey) IsValid() bool {
	if s.Revoked != nil {
		return false
	}

	if s.NotBefore != nil && time.Now().Before(*s.NotBefore) {
		return false
	}

	if s.NotAfter != nil && time.Now().After(*s.NotAfter) {
		return false
	}

	return true
}

func (s *SessionKey) MarshalBSON() ([]byte, error) {
	var err error
	var pub, priv []byte

	if s.PrivateKey != nil {
		priv, err = x509.MarshalECPrivateKey(s.PrivateKey)
		if err != nil {
			return nil, err
		}
	}

	// If there's a private key and a public key set then just save the private
	// key. The private key already contains a copy of the public key.
	if s.PublicKey != nil && s.PrivateKey == nil {
		pub, err = x509.MarshalPKIXPublicKey(s.PublicKey)
		if err != nil {
			return nil, err
		}
	}

	return bson.Marshal(struct {
		KeyId      string
		Revoked    *time.Time
		NotAfter   *time.Time
		NotBefore  *time.Time
		PublicKey  string
		PrivateKey string
	}{
		s.KeyId,
		s.Revoked, s.NotAfter, s.NotBefore,
		base64.StdEncoding.EncodeToString(pub),
		base64.StdEncoding.EncodeToString(priv),
	})
}

func (s *SessionKey) UnmarshalBSON(d []byte) error {
	v := struct {
		KeyId      string
		Revoked    *time.Time
		NotAfter   *time.Time
		NotBefore  *time.Time
		PublicKey  string
		PrivateKey string
	}{}
	if err := bson.Unmarshal(d, &v); err != nil {
		return err
	}

	s.KeyId = v.KeyId
	s.Revoked = v.Revoked
	s.NotAfter = v.NotAfter
	s.NotBefore = v.NotBefore

	if v.PrivateKey != "" {
		privb, err := base64.StdEncoding.DecodeString(v.PrivateKey)
		if err != nil {
			return err
		}

		priv, err := x509.ParseECPrivateKey(privb)
		if err != nil {
			return err
		}

		s.PrivateKey = priv
		s.PublicKey = priv.Public()
	}

	// If there was a private key then the public key was already set by
	// decoding that private key. No need to do this a second time (also it's
	// rather unlikely that both would be set).
	if v.PublicKey != "" && s.PublicKey == nil {
		pubb, err := base64.StdEncoding.DecodeString(v.PublicKey)
		if err != nil {
			return err
		}

		pubp, err := x509.ParsePKIXPublicKey(pubb)
		if err != nil {
			return err
		}

		pub, ok := pubp.(*ecdsa.PublicKey)
		if !ok {
			return fmt.Errorf("Failed to convert public key to *ecdsa.PublicKey")
		}

		s.PublicKey = pub
	}

	return nil
}