package aerotun

import (
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"sync"

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

// Version 2 binds data to fresh sender AND receiver boot identities. Directional
// keys and XChaCha nonces cannot overlap between peers or restarts.
const (
	headerSize    = 44 // magic(2), version(1), flags(1), sender(16), receiver(16), counter(8)
	flagHello     = 4
	flagAck       = 8
	flagConfirm   = 16
	maxPacketSize = 65535
)

type sessionID [16]byte

var zeroSession sessionID
var errReplay = errors.New("replayed or stale packet")
var errSession = errors.New("peer session is not established")

type replayWindow struct{ max, bits uint64 }

func (w *replayWindow) accept(n uint64) bool {
	if n == 0 {
		return false
	}
	if n > w.max {
		d := n - w.max
		if d >= 64 {
			w.bits = 1
		} else {
			w.bits = (w.bits << d) | 1
		}
		w.max = n
		return true
	}
	d := w.max - n
	if d >= 64 || w.bits&(uint64(1)<<d) != 0 {
		return false
	}
	w.bits |= uint64(1) << d
	return true
}

type protocolState struct {
	mu            sync.Mutex
	secret        [32]byte
	local, remote sessionID
	counter       uint64
	replay        replayWindow
	retired       map[sessionID]bool
	tx, rx        cipher.AEAD
}

func newProtocol(psk string) (*protocolState, error) {
	p := &protocolState{secret: sha256.Sum256([]byte(psk)), retired: make(map[sessionID]bool)}
	if _, err := rand.Read(p.local[:]); err != nil {
		return nil, err
	}
	return p, nil
}
func (p *protocolState) aead(sender, receiver sessionID) cipher.AEAD {
	mac := hmac.New(sha256.New, p.secret[:])
	mac.Write([]byte("AeroTUN-v2-directional-key"))
	mac.Write(sender[:])
	mac.Write(receiver[:])
	a, _ := chacha20poly1305.NewX(mac.Sum(nil))
	return a
}
func packetNonce(sender sessionID, counter uint64) []byte {
	nonce := make([]byte, chacha20poly1305.NonceSizeX)
	copy(nonce, sender[:])
	binary.BigEndian.PutUint64(nonce[16:], counter)
	return nonce
}
func (p *protocolState) seal(flags byte, payload []byte, target sessionID) ([]byte, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if target == zeroSession && flags != flagHello {
		target = p.remote
	}
	if target == zeroSession && flags != flagHello {
		return nil, errSession
	}
	if p.counter == ^uint64(0) {
		return nil, fmt.Errorf("packet counter exhausted; restart tunnel")
	}
	p.counter++
	b := make([]byte, headerSize, headerSize+len(payload)+chacha20poly1305.Overhead)
	binary.BigEndian.PutUint16(b, MagicHeader)
	b[2] = ProtocolVersion
	b[3] = flags
	copy(b[4:20], p.local[:])
	copy(b[20:36], target[:])
	binary.BigEndian.PutUint64(b[36:], p.counter)
	a := p.tx
	if target != p.remote || flags == flagHello || a == nil {
		a = p.aead(p.local, target)
	}
	return a.Seal(b, packetNonce(p.local, p.counter), payload, b), nil
}

// Authenticate before changing replay or peer state. Only an ACK addressed to
// our fresh boot ID (ACK or CONFIRM) establishes a session; HELLOs cannot
// replace a working peer.
func (p *protocolState) open(b []byte) (payload []byte, flags byte, sender sessionID, changed bool, err error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if len(b) < headerSize+chacha20poly1305.Overhead || binary.BigEndian.Uint16(b) != MagicHeader || b[2] != ProtocolVersion {
		err = errors.New("invalid AeroTUN v2 packet")
		return
	}
	flags = b[3]
	if flags != 0 && flags != FlagCompressedLZ4 && flags != FlagKeepalive && flags != flagHello && flags != flagAck && flags != flagConfirm {
		err = errors.New("invalid packet flags")
		return
	}
	copy(sender[:], b[4:20])
	var receiver sessionID
	copy(receiver[:], b[20:36])
	if sender == zeroSession || sender == p.local || (flags == flagHello && receiver != zeroSession) || (flags != flagHello && receiver != p.local) {
		err = errSession
		return
	}
	n := binary.BigEndian.Uint64(b[36:44])
	if n == 0 {
		err = errReplay
		return
	}
	a := p.rx
	if sender != p.remote || flags == flagHello || a == nil {
		a = p.aead(sender, receiver)
	}
	payload, err = a.Open(nil, packetNonce(sender, n), b[headerSize:], b[:headerSize])
	if err != nil {
		return
	}
	if flags != 0 && flags != FlagCompressedLZ4 && (len(payload) != 2 || binary.BigEndian.Uint16(payload) < 576 || binary.BigEndian.Uint16(payload) > 9000) {
		err = errSession
		return
	}
	if flags == flagHello {
		return
	}
	if sender != p.remote {
		if (flags != flagAck && flags != flagConfirm) || p.retired[sender] || len(p.retired) >= 1024 {
			err = errSession
			return
		}
		if p.remote != zeroSession {
			p.retired[p.remote] = true
		}
		p.remote = sender
		p.replay = replayWindow{}
		p.tx = p.aead(p.local, sender)
		p.rx = a
		changed = true
	}
	if !p.replay.accept(n) {
		err = errReplay
		return
	}
	return
}