// udpobfs -- simple point-to-point UDP obfuscation proxy
// Copyright (C) 2023-2024 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package udpobfs

import (
	"crypto/subtle"
	"io"
	"log"

	"golang.org/x/crypto/blowfish"
	"lukechampine.com/blake3"
)

const (
	App     = "go.cypherpunks.ru/udpobfs/v2"
	MACLen  = 6
	SeqLen  = 8
	SeedLen = 64
)

func MustNewBlowfish(key []byte) *blowfish.Cipher {
	ciph, err := blowfish.NewCipher(key)
	if err != nil {
		log.Fatal(err)
	}
	return ciph
}

type CryptoState struct {
	rxEnc, rxMAC   *blake3.Hasher
	txEnc, txMAC   *blake3.Hasher
	rxObfs, txObfs *blowfish.Cipher
	rxSeq, txSeq   [SeqLen]byte
}

func NewCryptoState(seed []byte, isInit bool) *CryptoState {
	var hInitEncKey, hInitMACKey [32]byte
	var hRespEncKey, hRespMACKey [32]byte
	var obfsInitKey, obfsRespKey [32]byte
	blake3.DeriveKey(hInitEncKey[:], App+" init enc", seed)
	blake3.DeriveKey(hInitMACKey[:], App+" init mac", seed)
	blake3.DeriveKey(hRespEncKey[:], App+" resp enc", seed)
	blake3.DeriveKey(hRespMACKey[:], App+" resp mac", seed)
	blake3.DeriveKey(obfsInitKey[:], App+" init obfs", seed)
	blake3.DeriveKey(obfsRespKey[:], App+" resp obfs", seed)
	hInitEnc := blake3.New(32, hInitEncKey[:])
	hInitMAC := blake3.New(MACLen, hInitMACKey[:])
	hRespEnc := blake3.New(32, hRespEncKey[:])
	hRespMAC := blake3.New(MACLen, hRespMACKey[:])
	obfsInit := MustNewBlowfish(obfsInitKey[:])
	obfsResp := MustNewBlowfish(obfsRespKey[:])
	state := CryptoState{}
	if isInit {
		state.rxEnc = hRespEnc
		state.rxMAC = hRespMAC
		state.rxObfs = obfsResp
		state.txEnc = hInitEnc
		state.txMAC = hInitMAC
		state.txObfs = obfsInit
	} else {
		state.rxEnc = hInitEnc
		state.rxMAC = hInitMAC
		state.rxObfs = obfsInit
		state.txEnc = hRespEnc
		state.txMAC = hRespMAC
		state.txObfs = obfsResp
	}
	return &state
}

func (state *CryptoState) Tx(dst, src []byte) []byte {
	if len(dst) != SeqLen+len(src) {
		log.Fatal("unexpected dst len")
	}
	if Incr(state.txSeq[MACLen:]) {
		Incr(state.txSeq[:MACLen])
	}
	copy(dst, state.txSeq[:])
	state.txEnc.Reset()
	MustWrite(state.txEnc, state.txSeq[:])
	if _, err := io.ReadFull(state.txEnc.XOF(), dst[SeqLen:]); err != nil {
		log.Fatal(err)
	}
	subtle.XORBytes(dst[SeqLen:], dst[SeqLen:], src)
	state.txMAC.Reset()
	MustWrite(state.txMAC, dst)
	state.txMAC.Sum(dst[:0])
	state.txObfs.Encrypt(dst[:SeqLen], dst[:SeqLen])
	return dst[:len(src)+SeqLen]
}

func (state *CryptoState) Rx(dst, src []byte) []byte {
	if len(dst) != len(src) {
		log.Fatal("unexpected dst len")
	}
	state.rxObfs.Decrypt(dst[:SeqLen], src[:SeqLen])
	var theirMAC [MACLen]byte
	copy(theirMAC[:], dst)
	copy(dst, state.rxSeq[:MACLen])
	state.rxMAC.Reset()
	MustWrite(state.rxMAC, dst[:SeqLen])
	MustWrite(state.rxMAC, src[SeqLen:])
	var ourMAC [MACLen]byte
	state.rxMAC.Sum(ourMAC[:0])
	if subtle.ConstantTimeCompare(ourMAC[:], theirMAC[:]) != 1 {
		Incr(dst[:MACLen])
		state.rxMAC.Reset()
		MustWrite(state.rxMAC, dst[:SeqLen])
		MustWrite(state.rxMAC, src[SeqLen:])
		state.rxMAC.Sum(ourMAC[:0])
		if subtle.ConstantTimeCompare(ourMAC[:], theirMAC[:]) != 1 {
			return nil
		}
	}
	copy(state.rxSeq[:], dst)
	state.rxEnc.Reset()
	MustWrite(state.rxEnc, dst[:SeqLen])
	if _, err := io.ReadFull(state.rxEnc.XOF(), dst); err != nil {
		log.Fatal(err)
	}
	subtle.XORBytes(dst, dst, src[SeqLen:])
	return dst[:len(src)-SeqLen]
}