package models

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"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 and from BSON and JSON. Other
// serializations can be added in the future as needed.
//
// The ExposePrivateKeysInJSON controls how JSON serialization of this struct
// works. When the field is set to false (the default) then serialization into
// JSON will never encode a private key, but may encode a public key. If this
// is set to true then the private key will be encoded into the JSON value and
// not the public key. SETTING THIS TO TRUE AND EXPOSING THE RESULTS TO THE
// USER IS A SECURITY ERROR so this should normally not be changed. This value
// of this field will never be persisted in any form.
//
// 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
	ExposePrivateKeysInJSON bool `json:"-" bson:"-"`
}

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
}

// MarshalPrivateKey marshals the private key to a X509 encoded base64 string
func (s *SessionKey) MarshalPrivateKey() (string, error) {
	priv, err := x509.MarshalECPrivateKey(s.PrivateKey)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(priv), nil
}

// MarshalPublicKey marshals the public key to an X509 encoded base64 string
func (s *SessionKey) MarshalPublicKey() (string, error) {
	pub, err := x509.MarshalPKIXPublicKey(s.PublicKey)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(pub), nil
}

// UnmarshalPrivateKey unmarshals the private key from a base64 encoded X509
// string into the public and private key fields.
func (s *SessionKey) UnmarshalPrivateKey(k string) error {
	privb, err := base64.StdEncoding.DecodeString(k)
	if err != nil {
		return err
	}

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

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

	return nil
}

// UnmarshalPublicKey unmarshals the public key from a base64 encoded X509
// string into the public key field.
func (s *SessionKey) UnmarshalPublicKey(k string) error {
	pubb, err := base64.StdEncoding.DecodeString(k)
	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
}

// MarshalJSON marshals a struct to JSON
//
// This method will have different behavior if the ExposePrivateKeysInJSON
// field is set in the struct (the default is false). If this field is set to
// true the private keys will be exposed in the JSON results. If it is false
// then private keys will not be exposed. The ExposePrivateKeysInJSON itself
// will never be serialized.
func (s *SessionKey) MarshalJSON() ([]byte, error) {
	var err error
	var privKey *string
	var pub, priv string

	if s.PrivateKey != nil && s.ExposePrivateKeysInJSON {
		priv, err = s.MarshalPrivateKey()
		if err != nil {
			return nil, err
		}
		if priv != "" {
			privKey = &priv
		}
	}

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

	return json.Marshal(struct {
		KeyId      string     `json:"key_id"`
		Revoked    *time.Time `json:"revoked,omitempty"`
		NotAfter   *time.Time `json:"not_after"`
		NotBefore  *time.Time `json:"not_before"`
		PublicKey  string     `json:"public_key"`
		PrivateKey *string    `json:"private_key,omitempty"`
	}{
		s.KeyId,
		s.Revoked, s.NotAfter, s.NotBefore,
		pub, privKey,
	})
}

// UnmarshalJSON unmarshals a struct from JSON.
//
// This method does attempt to unmarshal private keys.
func (s *SessionKey) UnmarshalJSON(d []byte) error {
	v := struct {
		KeyId      string     `json:"key_id"`
		Revoked    *time.Time `json:"revoked,omitempty"`
		NotAfter   *time.Time `json:"not_after"`
		NotBefore  *time.Time `json:"not_before"`
		PublicKey  string     `json:"public_key"`
		PrivateKey string     `json:"private_key"`
	}{}
	if err := json.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 != "" {
		if err := s.UnmarshalPrivateKey(v.PrivateKey); err != nil {
			return err
		}
	}

	// 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 {
		if err := s.UnmarshalPublicKey(v.PublicKey); err != nil {
			return err
		}
	}

	return nil
}

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

	if s.PrivateKey != nil {
		priv, err = s.MarshalPrivateKey()
		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 = s.MarshalPublicKey()
		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,
		pub, 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 != "" {
		if err := s.UnmarshalPrivateKey(v.PrivateKey); err != nil {
			return err
		}
	}

	// 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 {
		if err := s.UnmarshalPublicKey(v.PublicKey); err != nil {
			return err
		}
	}

	return nil
}