// Copyright 2024 The Forgejo Authors. All rights reserved.
// SPDX-License-Identifier: MIT

// Keying is a module that allows for subkeys to be deterministically generated
// from the same master key. It allows for domain separation to take place by
// using new keys for new subsystems/domains. These subkeys are provided with
// an API to encrypt and decrypt data. The module panics if a bad interaction
// happened, the panic should be seen as an non-recoverable error.
//
// HKDF (per RFC 5869) is used to derive new subkeys in a safe manner. It
// provides a KDF security property, which is required for Forgejo, as the
// secret key would be an ASCII string and isn't a random uniform bit string.
// XChaCha-Poly1305 (per draft-irtf-cfrg-xchacha-01) is used as AEAD to encrypt
// and decrypt messages. A new fresh random nonce is generated for every
// encryption. The nonce gets prepended to the ciphertext.
package keying

import (
	"bytes"
	"crypto/cipher"
	"crypto/hkdf"
	crand "crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"sync"
	"sync/atomic"

	"golang.org/x/crypto/chacha20poly1305"
)

// Specifies the context for which a subkey should be derived for.
var (
	// Used for the `push_mirror` table.
	PushMirror = deriveKey("pushmirror")
	// Used for the `two_factor` table.
	TOTP = deriveKey("totp")
	// Used for the `secret` table.
	ActionSecret = deriveKey("action_secret")
	// Used for the `task` table where type == TaskTypeMigrateRepo.
	MigrateTask = deriveKey("migrate_repo_task")
	// Used for the `webhook` table.
	Webhook = deriveKey("webhook")
	// Used for the `mirror` table.
	PullMirror = deriveKey("pullmirror")
)

var (
	// The hash used for HKDF.
	hash = sha256.New
	// The AEAD used for encryption/decryption.
	aead = chacha20poly1305.NewX
	// The pseudorandom key generated by HKDF-Extract.
	prk atomic.Value
)

// Set the main IKM for this module.
func Init(ikm []byte) {
	// Salt is intentionally left empty, it's not useful to Forgejo's use case.
	buf, err := hkdf.Extract(hash, ikm, nil)
	if err != nil {
		panic(err)
	}
	if ok := prk.CompareAndSwap(nil, buf); ok {
		return
	}
	// prk was already set
	old := prk.Load().([]byte)
	if bytes.Equal(old, buf) {
		return
	}
	panic("main IKM cannot be updated at runtime")
}

const (
	aeadKeySize   = chacha20poly1305.KeySize
	aeadNonceSize = chacha20poly1305.NonceSizeX
)

// Derive *the* key for a given context, this is a deterministic function.
// The same key will be provided for the same context.
func deriveKey(context string) Context {
	// wrap another sync.Once to prevent panic on initialization (prk would be nil)
	return Context{sync.OnceValue(func() cipher.AEAD {
		return expandPRK(prk.Load().([]byte), context)
	})}
}

func expandPRK(prk []byte, context string) cipher.AEAD {
	if len(prk) != sha256.Size {
		panic("keying: not initialized")
	}

	key, err := hkdf.Expand(hash, prk, context, aeadKeySize)
	if err != nil {
		panic(err)
	}
	e, err := aead(key)
	if err != nil {
		panic(err)
	}
	return e
}

type Context struct {
	aead func() cipher.AEAD
}

// Encrypts the specified plaintext with some additional data that is tied to
// this plaintext. The additional data can be seen as the context in which the
// data is being encrypted for, this is different than the context for which the
// key was derived; this allows for more granularity without deriving new keys.
// Avoid any user-generated data to be passed into the additional data. The most
// common usage of this would be to encrypt a database field, in that case use
// the ID and database column name as additional data. The additional data isn't
// appended to the ciphertext and may be publicly known, it must be available
// when decryping the ciphertext.
func (k Context) Encrypt(plaintext, additionalData []byte) []byte {
	nonce := make([]byte, aeadNonceSize)
	_, _ = crand.Read(nonce) // never returns an error

	// Returns the ciphertext of this plaintext.
	return k.aead().Seal(nonce, nonce, plaintext, additionalData)
}

// Decrypts the ciphertext and authenticates it against the given additional
// data that was given when it was encrypted. It returns an error if the
// authentication failed.
func (k Context) Decrypt(ciphertext, additionalData []byte) ([]byte, error) {
	if len(ciphertext) <= aeadNonceSize {
		return nil, errors.New("keying: ciphertext is too short")
	}

	nonce, ciphertext := ciphertext[:aeadNonceSize], ciphertext[aeadNonceSize:]

	return k.aead().Open(nil, nonce, ciphertext, additionalData)
}

// ColumnAndID generates a context that can be used as additional context for
// encrypting and decrypting data. It requires the column name and the row ID
// (this requires to be known beforehand). Be careful when using this, as the
// table name isn't part of this context. This means it's not bound to a
// particular table. The table should be part of the context that the key was
// derived for, in which case it binds through that.
func ColumnAndID(column string, id int64) []byte {
	return binary.BigEndian.AppendUint64(append([]byte(column), ':'), uint64(id))
}

// ColumnAndJSONSelectorAndID generates a context that can be used as additional context
// for encrypting and decrypting data. It requires the column name, JSON
// selector and the row ID (this requires to be known beforehand). Be careful
// when using this, as the table name isn't part of this context. This means
// it's not bound to a particular table. The table should be part of the context
// that the key was derived for, in which case it binds through that. Use this
// over `ColumnAndID` if you're encrypting data that's stored inside JSON.
// jsonSelector must be a unambiguous selector to the JSON field that stores the
// encrypted data.
func ColumnAndJSONSelectorAndID(column, jsonSelector string, id int64) []byte {
	return binary.BigEndian.AppendUint64(append(append([]byte(column), ':'), append([]byte(jsonSelector), ':')...), uint64(id))
}