// Copyright (C) 2020-2023 Sergey Matveev <stargrave@rnd.stcnet.ru>
// Use of this source code is governed by GPLv3 licence.

package tls

import (
	"crypto/cipher"
	"encoding/binary"

	"crypto/go.cypherpunks.ru/gogost/v5/gost34112012256"
	"crypto/go.cypherpunks.ru/gogost/v5/gost3412128"
	"crypto/go.cypherpunks.ru/gogost/v5/gost341264"
	"crypto/go.cypherpunks.ru/gogost/v5/mgm"
)

const GOSTAEADNonceSize = 8

type gostAEAD struct {
	nonceMask []byte
	tlsTree   *gost34112012256.TLSTree
	cipherer  func([]byte) cipher.Block
	aead      cipher.AEAD
}

func (g *gostAEAD) NonceSize() int        { return GOSTAEADNonceSize }
func (g *gostAEAD) Overhead() int         { return g.aead.Overhead() }
func (g *gostAEAD) explicitNonceLen() int { return 0 }

func (g *gostAEAD) keyRefresh(nonce []byte) {
	if len(nonce) != GOSTAEADNonceSize {
		panic("64-bit nonce expected")
	}
	seqNum := binary.BigEndian.Uint64(nonce)
	key, cached := g.tlsTree.DeriveCached(seqNum)
	if cached {
		return
	}
	cipher := g.cipherer(key)
	aead, err := mgm.NewMGM(cipher, cipher.BlockSize())
	if err != nil {
		panic(err)
	}
	g.aead = aead
}

func (g *gostAEAD) xorNonce(nonce []byte) {
	for i, b := range nonce {
		g.nonceMask[len(g.nonceMask)-GOSTAEADNonceSize+i] ^= b
	}
}

func (g *gostAEAD) Seal(out, nonce, plaintext, additionalData []byte) []byte {
	g.keyRefresh(nonce)
	g.xorNonce(nonce)
	ret := g.aead.Seal(out, g.nonceMask[:], plaintext, additionalData)
	g.xorNonce(nonce)
	return ret
}

func (g *gostAEAD) Open(out, nonce, ciphertext, additionalData []byte) ([]byte, error) {
	g.keyRefresh(nonce)
	g.xorNonce(nonce)
	ret, err := g.aead.Open(out, g.nonceMask[:], ciphertext, additionalData)
	g.xorNonce(nonce)
	return ret, err
}

func aeadGOST(
	blockSize int,
	tlsTreeParams gost34112012256.TLSTreeParams,
	cipherer func(key []byte) cipher.Block,
	key, nonceMask []byte,
) aead {
	if len(nonceMask) != blockSize {
		panic("tls: internal error: wrong nonce length")
	}
	ret := gostAEAD{
		nonceMask: make([]byte, blockSize),
		tlsTree:   gost34112012256.NewTLSTree(tlsTreeParams, key),
		cipherer:  cipherer,
	}
	copy(ret.nonceMask, nonceMask)
	ret.nonceMask[0] &= 0x7F
	ret.keyRefresh([]byte{0, 0, 0, 0, 0, 0, 0, 0})
	return &ret
}

func aeadGOSTR341112256WithKuznyechikMGML(key, nonceMask []byte) aead {
	return aeadGOST(
		gost3412128.BlockSize,
		gost34112012256.TLSGOSTR341112256WithKuznyechikMGML,
		func(key []byte) cipher.Block { return gost3412128.NewCipher(key) },
		key, nonceMask,
	)
}

func aeadGOSTR341112256WithMagmaMGML(key, nonceMask []byte) aead {
	return aeadGOST(
		gost341264.BlockSize,
		gost34112012256.TLSGOSTR341112256WithMagmaMGML,
		func(key []byte) cipher.Block { return gost341264.NewCipher(key) },
		key, nonceMask,
	)
}

func aeadGOSTR341112256WithKuznyechikMGMS(key, nonceMask []byte) aead {
	return aeadGOST(
		gost3412128.BlockSize,
		gost34112012256.TLSGOSTR341112256WithKuznyechikMGMS,
		func(key []byte) cipher.Block { return gost3412128.NewCipher(key) },
		key, nonceMask,
	)
}

func aeadGOSTR341112256WithMagmaMGMS(key, nonceMask []byte) aead {
	return aeadGOST(
		gost341264.BlockSize,
		gost34112012256.TLSGOSTR341112256WithMagmaMGMS,
		func(key []byte) cipher.Block { return gost341264.NewCipher(key) },
		key, nonceMask,
	)
}

var (
	gostSavedDefaultSignatureAlgorithms       []SignatureScheme
	gostSavedDefaultCipherSuites              []uint16
	gostSavedDefaultCipherSuitesNoAES         []uint16
	gostSavedCipherSuitesPreferenceOrder      []uint16
	gostSavedCipherSuitesPreferenceOrderNoAES []uint16
)

func GOSTInstallSignatureAlgorithms(sigs []SignatureScheme) {
	if gostSavedDefaultSignatureAlgorithms == nil {
		gostSavedDefaultSignatureAlgorithms = defaultSupportedSignatureAlgorithms
	}
	defaultSupportedSignatureAlgorithms = append(sigs, gostSavedDefaultSignatureAlgorithms...)
}

func GOSTInstallCipherSuites(suites []uint16) {
	if gostSavedDefaultCipherSuites == nil {
		gostSavedDefaultCipherSuites = defaultCipherSuitesTLS13
	}
	if gostSavedDefaultCipherSuitesNoAES == nil {
		gostSavedDefaultCipherSuitesNoAES = defaultCipherSuitesTLS13NoAES
	}
	if gostSavedCipherSuitesPreferenceOrder == nil {
		gostSavedCipherSuitesPreferenceOrder = cipherSuitesPreferenceOrder
	}
	if gostSavedCipherSuitesPreferenceOrderNoAES == nil {
		gostSavedCipherSuitesPreferenceOrderNoAES = cipherSuitesPreferenceOrderNoAES
	}
	defaultCipherSuitesTLS13 = append(suites, gostSavedDefaultCipherSuites...)
	defaultCipherSuitesTLS13NoAES = append(suites, gostSavedDefaultCipherSuitesNoAES...)
	cipherSuitesPreferenceOrder = append(suites, gostSavedCipherSuitesPreferenceOrder...)
	cipherSuitesPreferenceOrderNoAES = append(suites, gostSavedCipherSuitesPreferenceOrderNoAES...)
}

func GOSTInstall() {
	GOSTInstallSignatureAlgorithms([]SignatureScheme{
		GOSTR34102012256A,
		GOSTR34102012256B,
		GOSTR34102012256C,
		GOSTR34102012256D,
		GOSTR34102012512A,
		GOSTR34102012512B,
		GOSTR34102012512C,
	})
	GOSTInstallCipherSuites([]uint16{
		TLS_GOSTR341112_256_WITH_KUZNYECHIK_MGM_L,
		TLS_GOSTR341112_256_WITH_MAGMA_MGM_L,
		TLS_GOSTR341112_256_WITH_KUZNYECHIK_MGM_S,
		TLS_GOSTR341112_256_WITH_MAGMA_MGM_S,
	})
}