---

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"golang.org/x/crypto/argon2"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/mlkem"
	"lukechampine.com/blake3"
)

const (
	Magic0          = 0xAB
	Magic1          = 0xCD
	Version         = 0x02
	FixedPacketSize = 1280
	RekeyPackets    = 1 << 20
	RekeySeconds    = 600
	ServerPort      = "443"
	MaxFailures     = 3
	SwitchInterval  = 5 * time.Second
	CmdNewIP        = 0x03
	CmdNewSNI       = 0x04
	CmdRekey        = 0x05
)

var (
	RussianSNI = []string{
		"gosuslugi.ru",
		"vk.com",
		"ya.ru",
		"mail.ru",
		"sberbank.ru",
		"yandex.ru",
	}
	ServerIPPool = []string{
		"127.0.0.1",
	}
	ClientIPPool = []string{
		"127.0.0.1",
	}
	BridgeIP = "127.0.0.1"
	ExitIP   = "127.0.0.1"
)

type HandshakeKeys struct {
	X25519Priv []byte
	X25519Pub  []byte
	MLKEMPriv  *mlkem.DecapsulationKey768
	MLKEMPub   []byte
}

func GenerateHandshakeKeys() (*HandshakeKeys, error) {
	xPriv := make([]byte, 32)
	if _, err := rand.Read(xPriv); err != nil {
		return nil, err
	}
	xPub, err := curve25519.X25519(xPriv, curve25519.Basepoint)
	if err != nil {
		return nil, err
	}
	mPriv, err := mlkem.GenerateKey768()
	if err != nil {
		return nil, err
	}
	return &HandshakeKeys{
		X25519Priv: xPriv,
		X25519Pub:  xPub,
		MLKEMPriv:  mPriv,
		MLKEMPub:   mPriv.EncapsulationKey().Bytes(),
	}, nil
}

func DeriveKeyBLAKE3(shared, salt []byte, info string) []byte {
	h := blake3.New(32, nil)
	h.Write(shared)
	h.Write(salt)
	h.Write([]byte(info))
	return h.Sum(nil)
}

func DeriveKeyArgon2(password, salt []byte) []byte {
	return argon2.IDKey(password, salt, 2, 19*1024, 1, 32)
}

func RotateKeyBLAKE3(current []byte, counter uint32) []byte {
	salt := make([]byte, 4)
	binary.BigEndian.PutUint32(salt, counter)
	return DeriveKeyBLAKE3(current, salt, "XChangeNon-rotate")
}

func PadPayload(pt []byte) []byte {
	padded := make([]byte, FixedPacketSize)
	binary.BigEndian.PutUint16(padded[0:2], uint16(len(pt)))
	copy(padded[2:], pt)
	if _, err := rand.Read(padded[2+len(pt):]); err != nil {
		panic(err)
	}
	return padded
}

func UnpadPayload(p []byte) ([]byte, error) {
	if len(p) < 2 {
		return nil, fmt.Errorf("short")
	}
	n := binary.BigEndian.Uint16(p[0:2])
	if int(n) > len(p)-2 {
		return nil, fmt.Errorf("bad len")
	}
	return p[2 : 2+n], nil
}

func RandomSNI() string {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(RussianSNI))))
	return RussianSNI[n.Int64()]
}

func RandomServerIP() string {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(ServerIPPool))))
	return ServerIPPool[n.Int64()]
}

func Jitter(maxMs int) {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(maxMs)))
	time.Sleep(time.Duration(n.Int64()) * time.Millisecond)
}

type ReplayWindow struct {
	last uint64
	seen map[uint64]bool
	win  uint64
	mu   sync.Mutex
}

func NewReplayWindow() *ReplayWindow {
	return &ReplayWindow{seen: make(map[uint64]bool), win: 1024}
}

func (rw *ReplayWindow) Check(c uint64) bool {
	rw.mu.Lock()
	defer rw.mu.Unlock()
	if rw.last > rw.win && c <= rw.last-rw.win {
		return false
	}
	if rw.seen[c] {
		return false
	}
	rw.seen[c] = true
	if c > rw.last {
		rw.last = c
	}
	for k := range rw.seen {
		if k <= rw.last-rw.win {
			delete(rw.seen, k)
		}
	}
	return true
}

type DeadHand struct {
	mu        sync.Mutex
	probes    map[string]int
	lastProbe map[string]time.Time
	threshold int
	window    time.Duration
	onTrigger func(string)
}

func NewDeadHand(threshold int, window time.Duration, onTrigger func(string)) *DeadHand {
	return &DeadHand{
		probes:    make(map[string]int),
		lastProbe: make(map[string]time.Time),
		threshold: threshold,
		window:    window,
		onTrigger: onTrigger,
	}
}

func (dh *DeadHand) RecordProbe(ip string) {
	dh.mu.Lock()
	defer dh.mu.Unlock()
	now := time.Now()
	if last, ok := dh.lastProbe[ip]; ok && now.Sub(last) > dh.window {
		dh.probes[ip] = 0
	}
	dh.lastProbe[ip] = now
	dh.probes[ip]++
	if dh.probes[ip] >= dh.threshold {
		go dh.onTrigger(ip)
		dh.probes[ip] = 0
	}
}

func IsProbe(packet []byte) bool {
	if len(packet) < 44 {
		return true
	}
	if packet[0] != Magic0 || packet[1] != Magic1 {
		return true
	}
	return false
}

type AdaptiveRekey struct {
	mu             sync.Mutex
	packetCount    uint64
	suspicionLevel int
	lastRekey      time.Time
}

func NewAdaptiveRekey() *AdaptiveRekey {
	return &AdaptiveRekey{lastRekey: time.Now()}
}

func (ar *AdaptiveRekey) IncrementPacket() {
	ar.mu.Lock()
	defer ar.mu.Unlock()
	ar.packetCount++
}

func (ar *AdaptiveRekey) MarkSuspicious() {
	ar.mu.Lock()
	defer ar.mu.Unlock()
	ar.suspicionLevel++
}

func (ar *AdaptiveRekey) ShouldRekey() bool {
	ar.mu.Lock()
	defer ar.mu.Unlock()
	threshold := uint64(RekeyPackets)
	interval := time.Duration(RekeySeconds) * time.Second
	if ar.suspicionLevel > 0 {
		threshold = 1000
		interval = 30 * time.Second
	}
	if ar.packetCount >= threshold {
		return true
	}
	if time.Since(ar.lastRekey) >= interval {
		return true
	}
	return false
}

func (ar *AdaptiveRekey) OnRekey() {
	ar.mu.Lock()
	defer ar.mu.Unlock()
	ar.packetCount = 0
	ar.lastRekey = time.Now()
	if ar.suspicionLevel > 0 {
		ar.suspicionLevel--
	}
}

type XChangeNonState struct {
	mu          sync.Mutex
	Key         []byte
	SessionID   [8]byte
	SendCounter uint64
	RekeyCount  uint32
	CurrentIP   string
	CurrentSNI  string
	DeadHand    *DeadHand
	Rekeyer     *AdaptiveRekey
	Replay      *ReplayWindow
	CreatedAt   time.Time
	Conn        *net.UDPConn
	Addr        *net.UDPAddr
}

func NewXChangeNonState(key []byte, sid [8]byte, dh *DeadHand, ar *AdaptiveRekey) *XChangeNonState {
	return &XChangeNonState{
		Key:        key,
		SessionID:  sid,
		CurrentIP:  RandomServerIP(),
		CurrentSNI: RandomSNI(),
		DeadHand:   dh,
		Rekeyer:    ar,
		Replay:     NewReplayWindow(),
		CreatedAt:  time.Now(),
	}
}

func (s *XChangeNonState) RotateAll() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.RekeyCount++
	s.Key = RotateKeyBLAKE3(s.Key, s.RekeyCount)
	s.SendCounter = 0
	s.CurrentIP = RandomServerIP()
	s.CurrentSNI = RandomSNI()
	s.Replay = NewReplayWindow()
	s.CreatedAt = time.Now()
	s.Rekeyer.OnRekey()
}

func (s *XChangeNonState) Tick() {
	if s.Rekeyer.ShouldRekey() {
		s.RotateAll()
	}
}

func (s *XChangeNonState) Encrypt(pt []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	aead, err := chacha20poly1305.NewX(s.Key)
	if err != nil {
		return nil, err
	}
	padded := PadPayload(pt)
	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	aad := make([]byte, 0, 20)
	aad = append(aad, Magic0, Magic1, Version, 0x00)
	aad = append(aad, s.SessionID[:]...)
	cb := make([]byte, 8)
	binary.BigEndian.PutUint64(cb, s.SendCounter)
	aad = append(aad, cb...)
	ct := aead.Seal(nil, nonce, padded, aad)
	pkt := make([]byte, 0, 44+24+len(ct))
	pkt = append(pkt, Magic0, Magic1, Version, 0x00)
	pkt = append(pkt, s.SessionID[:]...)
	pkt = append(pkt, cb...)
	pkt = append(pkt, nonce...)
	pkt = append(pkt, ct...)
	s.SendCounter++
	s.Rekeyer.IncrementPacket()
	return pkt, nil
}

func (s *XChangeNonState) Decrypt(pkt []byte) ([]byte, error) {
	if len(pkt) < 68 {
		return nil, fmt.Errorf("short")
	}
	if pkt[0] != Magic0 || pkt[1] != Magic1 {
		return nil, fmt.Errorf("bad magic")
	}
	if pkt[2] != Version {
		return nil, fmt.Errorf("bad ver")
	}
	sid := pkt[4:12]
	c := binary.BigEndian.Uint64(pkt[12:20])
	nonce := pkt[20:44]
	ct := pkt[44:]
	if !s.Replay.Check(c) {
		return nil, fmt.Errorf("replay")
	}
	aead, err := chacha20poly1305.NewX(s.Key)
	if err != nil {
		return nil, err
	}
	aad := make([]byte, 0, 20)
	aad = append(aad, Magic0, Magic1, Version, 0x00)
	aad = append(aad, sid...)
	cb := make([]byte, 8)
	binary.BigEndian.PutUint64(cb, c)
	aad = append(aad, cb...)
	padded, err := aead.Open(nil, nonce, ct, aad)
	if err != nil {
		return nil, err
	}
	return UnpadPayload(padded)
}

func (s *XChangeNonState) NotifyClient(conn *net.UDPConn, addr *net.UDPAddr, cmd byte, data string) {
	payload := make([]byte, 0, 1+len(data))
	payload = append(payload, cmd)
	payload = append(payload, []byte(data)...)
	pkt, err := s.Encrypt(payload)
	if err != nil {
		return
	}
	Jitter(50)
	conn.WriteToUDP(pkt, addr)
}

func ServerHandshake(server *HandshakeKeys, clientXPub []byte) ([]byte, []byte, error) {
	xShared, err := curve25519.X25519(server.X25519Priv, clientXPub)
	if err != nil {
		return nil, nil, err
	}
	mShared, mCT := server.MLKEMPriv.EncapsulationKey().Encapsulate()
	return append(xShared, mShared...), mCT, nil
}

func ClientHandshake(client *HandshakeKeys, serverXPub, mlkemCT []byte) ([]byte, error) {
	xShared, err := curve25519.X25519(client.X25519Priv, serverXPub)
	if err != nil {
		return nil, err
	}
	mShared, err := client.MLKEMPriv.Decapsulate(mlkemCT)
	if err != nil {
		return nil, err
	}
	return append(xShared, mShared...), nil
}

type HostingAPI struct {
	Provider string
	Token    string
	BaseURL  string
}

type APIResponse struct {
	IP string `json:"ip"`
}

func (api *HostingAPI) RequestNewIP() (string, error) {
	reqBody, _ := json.Marshal(map[string]string{"action": "allocate_ip"})
	req, err := http.NewRequest("POST", api.BaseURL, bytes.NewBuffer(reqBody))
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+api.Token)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	var apiResp APIResponse
	if err := json.Unmarshal(body, &apiResp); err != nil {
		return "", err
	}
	return apiResp.IP, nil
}

func (api *HostingAPI) AddIPToPool(ip string) {
	ServerIPPool = append(ServerIPPool, ip)
	ClientIPPool = append(ClientIPPool, ip)
	fmt.Printf("[API] new IP added to pool: %s\n", ip)
}

type Chain struct {
	BridgeIP string
	ExitIP   string
	Enabled  bool
}

func (c *Chain) Forward(packet []byte) ([]byte, error) {
	if !c.Enabled {
		return packet, nil
	}
	exitAddr, err := net.ResolveUDPAddr("udp", c.ExitIP+":"+ServerPort)
	if err != nil {
		return nil, err
	}
	conn, err := net.DialUDP("udp", nil, exitAddr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	if _, err := conn.Write(packet); err != nil {
		return nil, err
	}
	buf := make([]byte, 4096)
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	return buf[:n], nil
}

func RunServer(server *HandshakeKeys, listenAddr string, api *HostingAPI, chain *Chain) error {
	addr, err := net.ResolveUDPAddr("udp", listenAddr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	fmt.Printf("[server] XChangeNon v0.2 listening on %s\n", listenAddr)

	dh := NewDeadHand(5, 10*time.Second, func(ip string) {
		fmt.Printf("[DEAD HAND] probe from %s — rotating\n", ip)
		if api != nil {
			newIP, err := api.RequestNewIP()
			if err == nil {
				api.AddIPToPool(newIP)
			}
		}
	})

	buf := make([]byte, 4096)
	for {
		n, clientAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			continue
		}
		if IsProbe(buf[:n]) {
			dh.RecordProbe(clientAddr.IP.String())
			continue
		}
		if n < 32 {
			continue
		}
		clientXPub := buf[:32]
		shared, mlkemCT, err := ServerHandshake(server, clientXPub)
		if err != nil {
			continue
		}
		salt := make([]byte, 32)
		rand.Read(salt)
		key := DeriveKeyBLAKE3(shared, salt, "XChangeNon-session")
		sid := [8]byte{}
		copy(sid[:], salt[:8])

		resp := make([]byte, 0, 32+len(mlkemCT)+32)
		resp = append(resp, server.X25519Pub...)
		resp = append(resp, mlkemCT...)
		resp = append(resp, salt...)

		Jitter(50)
		conn.WriteToUDP(resp, clientAddr)

		ar := NewAdaptiveRekey()
		state := NewXChangeNonState(key, sid, dh, ar)
		state.Conn = conn
		state.Addr = clientAddr
		go handleClient(state, chain)
	}
}

func handleClient(state *XChangeNonState, chain *Chain) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	go func() {
		for range ticker.C {
			state.Tick()
		}
	}()

	buf := make([]byte, 4096)
	for {
		state.Conn.SetReadDeadline(time.Now().Add(5 * time.Minute))
		n, _, err := state.Conn.ReadFromUDP(buf)
		if err != nil {
			return
		}
		pt, err := state.Decrypt(buf[:n])
		if err != nil {
			continue
		}
		if chain != nil && chain.Enabled {
			resp, err := chain.Forward(pt)
			if err != nil {
				continue
			}
			pt = resp
		}
		reply := append([]byte("echo: "), pt...)
		pkt, err := state.Encrypt(reply)
		if err != nil {
			continue
		}
		Jitter(50)
		state.Conn.WriteToUDP(pkt, state.Addr)
	}
}

type Client struct {
	mu           sync.Mutex
	CurrentIP    string
	CurrentIndex int
	Failures     int
	Conn         *net.UDPConn
	State        *XChangeNonState
}

func NewClient() *Client {
	return &Client{
		CurrentIndex: 0,
		CurrentIP:    ClientIPPool[0],
	}
}

func (c *Client) SwitchToIP(ip string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.Conn != nil {
		c.Conn.Close()
	}

	c.CurrentIP = ip
	fmt.Printf("[client] switching to IP: %s\n", ip)

	addr, err := net.ResolveUDPAddr("udp", ip+":"+ServerPort)
	if err != nil {
		return err
	}
	conn, err := net.DialUDP("udp", nil, addr)
	if err != nil {
		return err
	}
	c.Conn = conn
	c.Failures = 0
	return nil
}

func (c *Client) SwitchToNextIP() error {
	c.mu.Lock()
	c.CurrentIndex = (c.CurrentIndex + 1) % len(ClientIPPool)
	nextIP := ClientIPPool[c.CurrentIndex]
	c.mu.Unlock()
	return c.SwitchToIP(nextIP)
}

func (c *Client) MarkFailure() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.Failures++
}

func (c *Client) ShouldSwitch() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.Failures >= MaxFailures
}

func (c *Client) HandleCommand(pt []byte) {
	if len(pt) < 3 {
		return
	}
	if pt[0] != Magic0 || pt[1] != Magic1 {
		return
	}
	cmd := pt[2]
	data := string(pt[3:])
	switch cmd {
	case CmdNewIP:
		c.SwitchToIP(data)
	case CmdNewSNI:
		fmt.Printf("[client] new SNI: %s\n", data)
	case CmdRekey:
		fmt.Printf("[client] rekey requested\n")
	}
}

func (c *Client) Monitor() {
	ticker := time.NewTicker(SwitchInterval)
	defer ticker.Stop()
	for range ticker.C {
		if c.ShouldSwitch() {
			c.SwitchToNextIP()
		}
	}
}

func RunClient(serverAddr string) error {
	client, err := GenerateHandshakeKeys()
	if err != nil {
		return err
	}

	c := NewClient()
	if err := c.SwitchToIP(ClientIPPool[0]); err != nil {
		return err
	}
	defer c.Conn.Close()

	go c.Monitor()

	Jitter(50)
	if _, err := c.Conn.Write(client.X25519Pub); err != nil {
		return err
	}

	buf := make([]byte, 4096)
	c.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err := c.Conn.Read(buf)
	if err != nil {
		c.MarkFailure()
		return err
	}

	serverXPub := buf[:32]
	mlkemCT := buf[32 : 32+1088]
	salt := buf[32+1088 : 32+1088+32]

	shared, err := ClientHandshake(client, serverXPub, mlkemCT)
	if err != nil {
		return err
	}
	key := DeriveKeyBLAKE3(shared, salt, "XChangeNon-session")
	sid := [8]byte{}
	copy(sid[:], salt[:8])

	dh := NewDeadHand(5, 10*time.Second, func(string) {})
	ar := NewAdaptiveRekey()
	state := NewXChangeNonState(key, sid, dh, ar)
	c.State = state

	pkt, err := state.Encrypt([]byte("hello from client"))
	if err != nil {
		return err
	}
	Jitter(50)
	if _, err := c.Conn.Write(pkt); err != nil {
		c.MarkFailure()
		return err
	}

	c.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err = c.Conn.Read(buf)
	if err != nil {
		c.MarkFailure()
		return err
	}
	reply, err := state.Decrypt(buf[:n])
	if err != nil {
		return err
	}
	c.HandleCommand(reply)
	fmt.Printf("[client] reply: %s\n", reply)
	return nil
}

func main() {
	_ = DeriveKeyArgon2

	api := &HostingAPI{
		Provider: "selectel",
		Token:    "YOUR_API_TOKEN",
		BaseURL:  "https://api.selectel.ru/vpc/resell/v2/ips",
	}

	chain := &Chain{
		BridgeIP: BridgeIP,
		ExitIP:   ExitIP,
		Enabled:  false,
	}

	mode := "server"
	if mode == "server" {
		server, _ := GenerateHandshakeKeys()
		fmt.Printf("[server] ML-KEM pub: %d bytes\n", len(server.MLKEMPub))
		RunServer(server, ":"+ServerPort, api, chain)
	} else {
		RunClient("127.0.0.1:" + ServerPort)
	}
}