// TODO: somehow persist internal IPs
// TODO: key rotation
// TODO: dont trust claimed IP at all
// TODO: replay attacks are still a thing
package main

import (
	"crypto/hmac"
	"crypto/hpke"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"log"
	"net"
	"os"

	"baalvpn/shared"

	"github.com/cloudflare/circl/sign/mldsa/mldsa44"
)

type Peer struct {
	RealAddr     *net.UDPAddr
	InternalIP   string
	HexPublicKey string
	AuthSecret   []byte
}

const SERVER_PRIVKEY = "b137ffedfaf45cb4dc5339373cdc935fbcc7597dde71b86112d282c2bda4c252c95edb48c5bd92159c992dd601b3294e59fda2b605c176232b760e545f28d3cd8a6ba045d4ad630481521182c5aab0d045f68f12de2a6cdbaea9245576b715cf1329c19e90b009b33708da77f4836cbe0b722ee325e1fe01c5b12f62a86802621829329024600243115b8680c0929108a1285428240017710046819c46264a4825c8b425d818248b2411111204e28481a3248e40c671081988020280d93800180688e4a864e28030c9080c4c800158a46524c390db0426d3b22424c845d11810109211e4c270a348415234610a304d89202590a480c4006c91240e12352021b131014540db320c2104116342648c86681c8744d924089cc845e24002544231021109d48069431290490048e13466031421511068d980404c32909b404112826c8b16061a132e19295054280cd1c08d92444109188de34046db42441323460c3741e4144080b62154184618b744019520210860d1a86820a11014c04004320803c14184b86c0297852185200ac1508ca25194c69182b0715934260018450a115008b52d1009619a922c200430c89491c312520a312488a00d8b16521c13904434411b434504830110120410c2444c284210b22003338de0228e01306448b60844060c19294618830000832400a780a1469240422440a864d1184dd9c084cb108a08363090464e049364899670a3b2400c018820282988b81018b7891c39465a84700c134811b48dda464981c20103044cc30045d9320de44880e226822239214c840c04274213c8251b08910205202116299b380400138424046808352219411091306a528080a22686c1444ed3002ecb18621a1410daa8205290315928661a37911b188889002014242a419268204189591641d1b664d128054c2664030645cab26484a0690a478c19b5812282695132409c344610b448ccc80518a545224489a48229d0a80d214249c4a66dda8228d4022d9b0441c820119a88215b1029944484d1b29110c90d14125013c4818bc228921070514070c922698a2890014620db98702334020b468d21014912a4901325519906095c3460dcb0911c37418a422dc3224a2211698110401282810b4670a3060c61a64843427053404ca2a68048486440246153966c90b86d1bb50863a029d0384964422d10258c59a285c8c89193860823b748032185d0b810a428524ca64c1a29490284815b008c0ca969c1a6409c36b1659c04f7e6f1cc28e6944fafdf06980608ec937a242945a9949704a47ae87cb6aa3b6d85d65f3f9cbcd301505fe0b19607d67f82ad539580c667927fe53c13af2e1a5177fd6f6f8a30829b363f586b92e5e9979c122393a0415e5b29fcf9da93196c10cad9c5038dd227961bc42f3f663d3647f9b811ff9ccaaf17d0a59b1fed877448412c254e3e7936023d4ef9b6b489080eaa04ca14e90e6947ff1d6320fc37acfa3d44b2710fda267e824513a056d7f5666bfac613008b1fba4f2fa12316b16ef3258ff9df5e1301835e76537318626d1d995637892b81d75ab7ed59ac0fd94af24bae94f457c88645b7de7f87faff893bda4bc55a54205b34eea275d02503eb3275b200f6b5fc2cdd82a2f46027e29b6a070c0dbe008933943861ed5e4566e7d1757084ef614f58b4237211dda05ef06bacf6ddc67d8f293fc874cf8eaaba7f4cea273e23fddfcbd40ee1ea7cf2bba94b35cded5682db8e451f5596082775055200da0c4de55834aa5e330e8d797d0d882243cca9eaa04452cbfdc50ebf5bdbacc1eee49c679eac90d13867b91e95a15a9a822ce61691f8a7aecd438d0d701534332688a7047734cc5c9ed538f7f8476c1fb8cdb44eb6db972815dc60464eb1931e1b8102838fb83b6a26458381001a70a61a7f68dd0ddbb077c5a86b48badac51271b01d5efbf019053bf6911ae8bf18fd77e70a2d0288d87af5a6035d52d3328262d72b11e1a597911ad7ec2349c106b25fd0207c31bd6f261104c1ec81f662031def52edce9eafb5cd30a4a77d1d3f4c030dcbcd13a110e1934c91323d2e8bcf57a4bd0f5f86fdd6f322af7791f65fa1af298ef6069863057847758c9285a214c0d264070a870c24a8d0a859d704a89cd0bd448c2175540259d92bae09ccf8e1ef5ba4f2ab8e409ab2027f17103ebe734d11984c8c3d4f337036cb3066413bcd488426ab17128829e88a73171d7f507e3405ba2b5b5bba6334cd51bc8861eff8443dad08670ff91bf45e918c832bfae923daae2c5425f032372257ba12d1b70d4eb38e96707270b5b3b5dff36c4f39526adfcfd96ec806828f431230d4637147e5a20ed646c6e5d81b759db34bde3e7fd39e2b16b774c5d8aa960df7cac49ba1b456ee9f3a0c8a177a6aea0f5d639e9efda6976da073b87b5e02d1a63da02d3b2e5710d4f2b76c7f01c747f328b3ccee5daf30a1a780ecd8896d376082468f1ec08b46a1533a6681c05b02d4a66327728d9014a3f81431f49b3aa82936d8224827606ef7faf12e4975832abd63513a02ffd72499cb33d305dcf1fe37006f156b04a02c3973c5e014f4733c1539f84921e7e1da5ee5d5cca7eaf2462bb4fdea2c91affea2eedfd96b6e706e718d1ff5b4a42b04a8fd20eade15c604bacd169508f3b06143e27e8bf817bf5f1fbd56a8a900ea45eb0c0c62447f9f964ae7b63d354270d02d58fa45b1f461fdceb7fa6bae89798ff7358de44adb58ca19aa989a34b3bb165c1b39054a7e2c1ec67d2fe7f80fc929f70116c41463e85b0e666709652a635eb6b2f1fc021e2e2903631a699b05015e905c68bcfc66dad19e87cb1c4dfb513d2bc473e72410cb3028cfb80295adde6aaa86737bea8b2ad135941d3ef2e17f0de620bd80e3d2aab1272bb7a44f86bcc5ea234910ab3cc15a300c48019d5035f6caf06b9602cc238fb1061a77067464ff0c4f0fed3d8abfccaed26f6e3945e6ec8c31dc92e9f04d1af749705af3612c73ed33e2f93f53eef4c29d48fd66cdb92cfadb6167ede617e71beff15cd231fefb95dfe9affcad1e67086474bfc49a3d3bdc34dd2017a3fc4d45bb2324cd31df6086efe84bc58891ce54bb24ca390297d0600ffc9e112c22860b58a32c21e5961538e30f835b2b0043579441877d858d1d44e1dc9902438bc3a1af68a52de945d7b857587bf37f85992f67e2388380bc3e3387262f9f1dde51491f7fee5202eba070bf82caa395ffeed4eb435073cdd52f40f48259e3c274588820359f4275dd0a097e5faae86aa29c399953b0e20261e3360f7a539d4eb60ccf2c677441e1215ec3fc0777ba613464ec48ea33238606d15e9c544d60d87aa39bd0b80f24a27fea01872eb18a0b38b374d609a14b057a8c6abba736e16ab9fb8a941df040883f9e0270a40e483664168f6c6f57071d06215b081313a221e75818ad80bdf59c84cf59f75af3026c190782b9a386e3c46479561ab413569fc1ca198f39b2b77fb9f21fd9de72998e5183f3006ac0ad36cdf3568aef4346608dc6a05f5cb7fed731cc731796704ed5c009e8fa7749625f8bdd78c28332104bcb9f8b73b786ef36472557f2231ca4ae"

var (
	conn                *net.UDPConn
	multicastKey        = make([]byte, 32)
	authSecretsByPubKey = shared.NewTMap[string, []byte]()

	peers           = shared.NewTMap[string, *Peer]()
	peersByInternal = shared.NewTMap[string, *Peer]()
	peersByPubKey   = shared.NewTMap[string, *Peer]()
)

func main() {
	if len(os.Args) > 1 && os.Args[1] == "keygen" {
		pubKey, privKey, err := mldsa44.GenerateKey(rand.Reader)
		if err != nil {
			panic(err) // cant continue
		}
		fmt.Println("pubKey: " + hex.EncodeToString(pubKey.Bytes()))
		fmt.Println("privKey: " + hex.EncodeToString(privKey.Bytes()))
		return
	}

	if os.Geteuid() != 0 {
		panic("root permissions needed")
	}

	if _, err := rand.Read(multicastKey); err != nil {
		panic(err) // should never happen
	}

	receivePackets()
}

func receivePackets() {
	log.Println("Listening for packets...")

	listenAddr, err := net.ResolveUDPAddr("udp", ":38000")
	if err != nil {
		panic(err) // cant continue
	}
	conn, err = net.ListenUDP("udp", listenAddr)
	if err != nil {
		panic(err) // cant continue
	}
	defer conn.Close()

	buffer := make([]byte, 50000)
	for {
		n, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Println(err)
			continue
		}
		if n == 0 {
			continue
		}

		req := make([]byte, n)
		copy(req, buffer[:n])
		go handleReq(req, addr)
	}
}

func handleReq(req []byte, addr *net.UDPAddr) {
	if len(req) <= 4 {
		log.Println("packet too short")
		return
	}
	version := binary.LittleEndian.Uint16(req[0:])
	if version != shared.PROTO_VERSION {
		log.Println("mismatched packet version")
		return
	}

	pktType := binary.LittleEndian.Uint16(req[2:])

	switch pktType {
	case shared.REQ_GET_CHALLENGE:
		pubKey, err := hpke.MLKEM768X25519().NewPublicKey(req[4:])
		if err != nil {
			log.Println("invalid pubkey")
			return
		}

		hexPubKey := hex.EncodeToString(pubKey.Bytes())

		authSecret := make([]byte, 32)
		if _, err = rand.Read(authSecret); err != nil {
			panic(err) // should never happen
		}

		authSecretsByPubKey.Set(hexPubKey, authSecret)

		ciphertext, err := hpke.Seal(pubKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn-challenge"), authSecret)
		if err != nil {
			log.Println("failed to Seal the authSecret:", err)
			return
		}

		sendTo(addr, buildSignedPkt(shared.RESP_GET_CHALLENGE, ciphertext))
	case shared.REQ_REGISTER:
		if len(req) <= 36 {
			log.Println("REGISTER packet too short")
			return
		}
		auth := req[4:36]

		pubKey, err := hpke.MLKEM768X25519().NewPublicKey(req[36:])
		if err != nil {
			log.Println("invalid pubkey")
			return
		}

		hexPubKey := hex.EncodeToString(pubKey.Bytes())

		authSecret, ok := authSecretsByPubKey.GetOK(hexPubKey)
		if !ok {
			log.Println("unknown pubkey tried to REGISTER")
			return
		}
		defer authSecretsByPubKey.Delete(hexPubKey)

		mac := hmac.New(sha256.New, authSecret)
		mac.Write(pubKey.Bytes())
		if !hmac.Equal(mac.Sum(nil), auth) {
			log.Println("failed to authenticate REGISTER")
			return
		}

		var internalIP string
		if peer, ok := peersByPubKey.GetOK(hexPubKey); ok {
			internalIP = peer.InternalIP
			peer.AuthSecret = authSecret
			peers.Delete(peer.RealAddr.String())
			peer.RealAddr = addr
			peers.Set(addr.String(), peer)
		} else {
			internalIP = randomIP().String()
			peer := &Peer{
				RealAddr:     addr,
				InternalIP:   internalIP,
				HexPublicKey: hexPubKey,
				AuthSecret:   authSecret,
			}
			peers.Set(addr.String(), peer)
			peersByInternal.Set(internalIP, peer)
			peersByPubKey.Set(hexPubKey, peer)
		}

		ciphertext, err := hpke.Seal(pubKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn-register"), multicastKey)
		if err != nil {
			log.Println("failed to Seal the multicastKey:", err)
			return
		}

		sendTo(addr, buildSignedPkt(shared.RESP_REGISTER, net.ParseIP(internalIP), ciphertext))
	case shared.UNICAST_PKT:
		if len(req) < 68 {
			log.Println("UNICAST_PKT packet too short")
			return
		}
		peer := peers.Get(addr.String())
		if peer == nil {
			log.Println("UNICAST_PKT from an unregistered peer")
			return
		}

		mac := hmac.New(sha256.New, peer.AuthSecret)
		mac.Write(req[36:])
		if !hmac.Equal(mac.Sum(nil), req[4:36]) {
			log.Println("failed to authenticate UNICAST_PKT")
			return
		}

		destIP := net.IP(req[36:52]).String()
		srcIP := net.IP(req[52:68]).String()
		if srcIP != peer.InternalIP {
			log.Println("rejected spoofed srcIP in UNICAST_PKT")
			return
		}

		if destPeer := peersByInternal.Get(destIP); destPeer != nil {
			log.Println(srcIP + " -> " + destIP)
			sendTo(destPeer.RealAddr, req)
		} else {
			log.Println(srcIP + " -/> " + destIP + " (unrecognized IP)")
		}
	case shared.BROADCAST_PKT:
		if len(req) <= 36 {
			log.Println("BROADCAST_PKT packet too short")
			return
		}
		peer := peers.Get(addr.String())
		if peer == nil {
			log.Println("BROADCAST_PKT from unregistered peer:", addr.String())
			return
		}

		mac := hmac.New(sha256.New, peer.AuthSecret)
		mac.Write(req[36:])
		if !hmac.Equal(mac.Sum(nil), req[4:36]) {
			log.Println("failed to authenticate BROADCAST_PKT")
			return
		}

		log.Println(peer.InternalIP + " -> *")
		broadcast(req, peer.InternalIP)
	case shared.REQ_GET_PUBKEY:
		if len(req) < 52 {
			log.Println("REQ_GET_PUBKEY packet too short")
			return
		}
		peerIP := net.IP(req[36:52]).String()
		peer := peersByInternal.Get(peerIP)
		if peer != nil {
			mac := hmac.New(sha256.New, peer.AuthSecret)
			mac.Write(req[36:])
			if !hmac.Equal(mac.Sum(nil), req[4:36]) {
				log.Println("failed to authenticate REQ_GET_PUBKEY")
				return
			}

			pubKey, err := hex.DecodeString(peer.HexPublicKey)
			if err != nil {
				panic(err) // should never happen
			}
			sendTo(addr, buildSignedPkt(shared.RESP_GET_PUBKEY, net.ParseIP(peerIP), pubKey))
		}
	case shared.REQ_ESTABLISH, shared.RESP_ESTABLISH:
		if len(req) < 68 {
			log.Println("ESTABLISH packet too short")
			return
		}
		peer := peers.Get(addr.String())
		if peer == nil {
			log.Println("ESTABLISH from unregistered peer:", addr.String())
			return
		}

		mac := hmac.New(sha256.New, peer.AuthSecret)
		mac.Write(req[36:])
		if !hmac.Equal(mac.Sum(nil), req[4:36]) {
			log.Println("failed to authenticate ESTABLISH")
			return
		}

		srcIP := net.IP(req[52:68]).String()
		if srcIP != peer.InternalIP {
			log.Println("rejected spoofed srcIP in ESTABLISH")
			return
		}

		dstIP := net.IP(req[36:52]).String()
		dstPeer := peersByInternal.Get(dstIP)
		if dstPeer == nil {
			log.Println("tried to ESTABLISH with an unknown peer")
			return
		}
		sendTo(dstPeer.RealAddr, req)
	default:
		log.Println("unknown packet type")
	}
}

func buildSignedPkt(pktType uint16, parts ...[]byte) []byte {
	pkt := shared.BuildPkt(pktType, parts...)

	privKeyBytes, err := hex.DecodeString(SERVER_PRIVKEY)
	if err != nil {
		panic(err) // fatal misconfiguration
	}
	var privKey mldsa44.PrivateKey
	if err := privKey.UnmarshalBinary(privKeyBytes); err != nil {
		panic(err) // fatal misconfiguration
	}

	sig, err := privKey.Sign(rand.Reader, pkt[4:], nil)
	if err != nil {
		panic(err) // should never happen
	}

	return append(pkt[:4], append(sig, pkt[4:]...)...)
}

func broadcast(pkt []byte, except string) {
	peers.ForEach(func(k string, v *Peer) {
		if v.InternalIP == except {
			return
		}
		go sendTo(v.RealAddr, pkt)
	})
}

func sendTo(addr *net.UDPAddr, pkt []byte) {
	if _, err := conn.WriteToUDP(pkt, addr); err != nil {
		log.Println(err)
	}
}

// fd00:baa1::/32
func randomIP() net.IP {
	ip := make(net.IP, 16)
	ip[0] = 0xfd
	ip[1] = 0x00
	ip[2] = 0xba
	ip[3] = 0xa1
	if _, err := rand.Read(ip[4:]); err != nil {
		panic(err) // should never happen
	}
	return ip
}