// Package push sends Web Push messages (RFC 8030), encrypted for the browser (RFC 8291) and signed
// for the push service with VAPID (RFC 8292). Only the standard library is used.
package push

import (
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

var b64 = base64.RawURLEncoding

// decode64 accepts base64url with or without padding, which is how browsers hand out keys.
func decode64(s string) ([]byte, error) {
	return b64.DecodeString(strings.TrimRight(s, "="))
}

// Keys is the server's VAPID key pair. Browsers are given the public key when they subscribe, and
// the push service only accepts messages signed with the matching private key.
type Keys struct {
	priv *ecdsa.PrivateKey
}

// GenerateKeys makes a new VAPID key pair.
func GenerateKeys() (*Keys, error) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	return &Keys{priv: priv}, nil
}

// ParseKeys reads a private key written by Keys.Private.
func ParseKeys(private string) (*Keys, error) {
	raw, err := decode64(private)
	if err != nil {
		return nil, fmt.Errorf("VAPID key: %w", err)
	}
	priv, err := ecdsa.ParseRawPrivateKey(elliptic.P256(), raw)
	if err != nil {
		return nil, fmt.Errorf("VAPID key: %w", err)
	}
	return &Keys{priv: priv}, nil
}

// Private is the private key as base64url, for storing.
func (k *Keys) Private() string {
	raw, _ := k.priv.Bytes()
	return b64.EncodeToString(raw)
}

// Public is the uncompressed public key as base64url: the applicationServerKey for subscribing.
func (k *Keys) Public() string {
	raw, _ := k.priv.PublicKey.Bytes()
	return b64.EncodeToString(raw)
}

// vapid returns the Authorization header for a push service: a JWT naming its origin, signed ES256.
func (k *Keys) vapid(endpoint, subject string, now time.Time) (string, error) {
	u, err := url.Parse(endpoint)
	if err != nil {
		return "", err
	}
	claims, _ := json.Marshal(map[string]any{
		"aud": u.Scheme + "://" + u.Host,
		"exp": now.Add(12 * time.Hour).Unix(),
		"sub": subject,
	})
	unsigned := b64.EncodeToString([]byte(`{"typ":"JWT","alg":"ES256"}`)) + "." + b64.EncodeToString(claims)
	digest := sha256.Sum256([]byte(unsigned))
	r, s, err := ecdsa.Sign(rand.Reader, k.priv, digest[:])
	if err != nil {
		return "", err
	}
	sig := make([]byte, 64)
	r.FillBytes(sig[:32])
	s.FillBytes(sig[32:])
	return "vapid t=" + unsigned + "." + b64.EncodeToString(sig) + ", k=" + k.Public(), nil
}

// Subscription is what a browser's PushManager.subscribe() returns.
type Subscription struct {
	Endpoint string
	P256dh   string // the browser's public key, base64url
	Auth     string // its 16-byte auth secret, base64url
}

// Check validates the keys, so a bad subscription is refused when it's saved rather than on send.
func (s Subscription) Check() error {
	pub, err := decode64(s.P256dh)
	if err != nil || len(pub) != 65 {
		return errors.New("the browser's push key isn't valid")
	}
	if _, err := ecdh.P256().NewPublicKey(pub); err != nil {
		return errors.New("the browser's push key isn't valid")
	}
	auth, err := decode64(s.Auth)
	if err != nil || len(auth) != 16 {
		return errors.New("the browser's push secret isn't valid")
	}
	return nil
}

const recordSize = 4096

// encrypt is RFC 8291's aes128gcm encoding of one message for one browser. The salt and the
// sender's one-off key are passed in so the RFC's test vector can check it.
func encrypt(plaintext, uaPublic, authSecret []byte, as *ecdh.PrivateKey, salt []byte) ([]byte, error) {
	ua, err := ecdh.P256().NewPublicKey(uaPublic)
	if err != nil {
		return nil, err
	}
	secret, err := as.ECDH(ua)
	if err != nil {
		return nil, err
	}
	asPublic := as.PublicKey().Bytes()

	prkKey, err := hkdf.Extract(sha256.New, secret, authSecret)
	if err != nil {
		return nil, err
	}
	keyInfo := "WebPush: info\x00" + string(uaPublic) + string(asPublic)
	ikm, err := hkdf.Expand(sha256.New, prkKey, keyInfo, 32)
	if err != nil {
		return nil, err
	}
	prk, err := hkdf.Extract(sha256.New, ikm, salt)
	if err != nil {
		return nil, err
	}
	cek, err := hkdf.Expand(sha256.New, prk, "Content-Encoding: aes128gcm\x00", 16)
	if err != nil {
		return nil, err
	}
	nonce, err := hkdf.Expand(sha256.New, prk, "Content-Encoding: nonce\x00", 12)
	if err != nil {
		return nil, err
	}

	block, err := aes.NewCipher(cek)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	// One record, ended with the last-record delimiter and no padding.
	padded := append(append([]byte{}, plaintext...), 0x02)

	header := make([]byte, 0, 16+4+1+len(asPublic))
	header = append(header, salt...)
	header = binary.BigEndian.AppendUint32(header, recordSize)
	header = append(header, byte(len(asPublic)))
	header = append(header, asPublic...)
	return gcm.Seal(header, nonce, padded, nil), nil
}

// MaxPayload is the most a message can carry. Push services take 4096 bytes after encryption.
const MaxPayload = recordSize - 86 - 16 - 1

// Message is one push.
type Message struct {
	Payload []byte
	// TTL is how long the push service keeps it for a phone that's offline.
	TTL time.Duration
	// Urgency is very-low, low, normal or high. High wakes a phone in battery saving.
	Urgency string
	// Topic replaces an undelivered message with the same topic (at most 32 base64url characters).
	Topic string
}

// Sender delivers messages. Subject is a mailto: or https: URL the push service can contact.
type Sender struct {
	Keys    *Keys
	Subject string
	Client  *http.Client
}

// ErrGone means the browser has unsubscribed or the subscription expired: delete it.
var ErrGone = errors.New("push subscription is gone")

// Send encrypts and posts one message to one subscription.
func (s *Sender) Send(ctx context.Context, sub Subscription, m Message) error {
	if len(m.Payload) > MaxPayload {
		return fmt.Errorf("push payload is %d bytes, over %d", len(m.Payload), MaxPayload)
	}
	uaPublic, err := decode64(sub.P256dh)
	if err != nil {
		return err
	}
	auth, err := decode64(sub.Auth)
	if err != nil {
		return err
	}
	as, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return err
	}
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return err
	}
	body, err := encrypt(m.Payload, uaPublic, auth, as, salt)
	if err != nil {
		return err
	}
	authz, err := s.Keys.vapid(sub.Endpoint, s.Subject, time.Now())
	if err != nil {
		return err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, sub.Endpoint, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("Content-Encoding", "aes128gcm")
	req.Header.Set("Authorization", authz)
	req.Header.Set("TTL", strconv.Itoa(int(m.TTL.Seconds())))
	if m.Urgency != "" {
		req.Header.Set("Urgency", m.Urgency)
	}
	if m.Topic != "" {
		req.Header.Set("Topic", m.Topic)
	}
	client := s.Client
	if client == nil {
		client = http.DefaultClient
	}
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	msg, _ := io.ReadAll(io.LimitReader(res.Body, 512))
	switch {
	case res.StatusCode == http.StatusNotFound || res.StatusCode == http.StatusGone:
		return ErrGone
	case res.StatusCode >= 300:
		return fmt.Errorf("push service answered %s: %s", res.Status, strings.TrimSpace(string(msg)))
	}
	return nil
}

// pushHosts are the push services of the browsers in use. Subscriptions are only accepted for
// these, so the panel can't be pointed at an arbitrary URL.
var pushHosts = []string{
	"fcm.googleapis.com",        // Chrome, Edge on Android, Samsung Internet, Brave, Opera
	"push.services.mozilla.com", // Firefox
	"web.push.apple.com",        // Safari
	"notify.windows.com",        // Edge on Windows
}

// AllowedEndpoint reports whether an endpoint is an https URL on a known push service.
func AllowedEndpoint(endpoint string) bool {
	u, err := url.Parse(endpoint)
	if err != nil || u.Scheme != "https" || u.User != nil || u.Port() != "" {
		return false
	}
	host := strings.ToLower(u.Hostname())
	for _, h := range pushHosts {
		if host == h || strings.HasSuffix(host, "."+h) {
			return true
		}
	}
	return false
}