// TODO: slice bounds checking
// TODO: authentication
// TODO: somehow persist IPs
// TODO: handle multiple peers behind one NAT
// TODO: key rotation
// TODO: sha256 -> HKDF
// TODO: recovery
package main

import (
	"crypto/rand"
	"encoding/binary"
	"log"
	"net"
	"os"
	"sync"

	"baalvpn/shared"
)

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

var (
	peers      = map[string]*Peer{}
	peersMutex sync.Mutex
	conn       *net.UDPConn
)

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

	receivePackets()
}

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

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

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

		go handleReq(buffer[:n], 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_REGISTER:
		peersMutex.Lock()
		defer peersMutex.Unlock()

		pubKey := req[4:]

		internalIP := randomIP().String()

		if destPeer, ok := peers[addr.IP.String()]; ok {
			internalIP = destPeer.InternalIP
		} else {
			peers[addr.IP.String()] = &Peer{
				RealAddr:   addr,
				InternalIP: internalIP,
				PublicKey:  pubKey,
			}
		}

		resp := shared.BuildPkt(shared.RESP_REGISTER, net.ParseIP(internalIP))
		sendTo(addr, resp)
	case shared.ENC_PKT:
		peersMutex.Lock()
		peer := peers[addr.IP.String()]
		peersMutex.Unlock()
		if peer == nil {
			log.Println("data from unregistered peer:", addr.String())
			return
		}

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

		if destPeer := getPeerByInternalIP(destIP); destPeer != nil {
			log.Println(srcIP + " -> " + destIP)
			sendTo(destPeer.RealAddr, req)
		} else {
			log.Println(srcIP + " -/> " + destIP + " (unrecognized IP)")
		}
	case shared.BROADCAST_PKT:
		peersMutex.Lock()
		peer := peers[addr.IP.String()]
		peersMutex.Unlock()
		if peer == nil {
			log.Println("data from unregistered peer:", addr.String())
			return
		}

		destIP := net.IP(req[42:58]).String()
		log.Println(peer.InternalIP + " -> " + destIP)
		broadcast(req, peer.InternalIP)
	case shared.REQ_GET_PUBKEY:
		peerIP := net.IP(req[4:20]).String()
		peer := getPeerByInternalIP(peerIP)
		if peer == nil {
			log.Println(addr.String(), "tried to get pubkey of unknown peer", peerIP)
			return
		}

		resp := shared.BuildPkt(shared.RESP_GET_PUBKEY, append(net.ParseIP(peer.InternalIP), peer.PublicKey...))
		sendTo(addr, resp)
	case shared.REQ_ESTABLISH, shared.RESP_ESTABLISH:
		peersMutex.Lock()
		peer := peers[addr.IP.String()]
		peersMutex.Unlock()

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

		dstIP := net.IP(req[4:20]).String()
		dstPeer := getPeerByInternalIP(dstIP)
		sendTo(dstPeer.RealAddr, req)
	default:
		log.Println("unknown packet type")
	}
}

func broadcast(pkt []byte, except string) {
	peersMutex.Lock()
	defer peersMutex.Unlock()
	for _, p := range peers {
		if p.InternalIP == except {
			continue
		}
		sendTo(p.RealAddr, pkt)
	}
}

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

func getPeerByInternalIP(internalIP string) *Peer {
	peersMutex.Lock()
	defer peersMutex.Unlock()
	for _, p := range peers {
		if p.InternalIP == internalIP {
			return p
		}
	}
	return nil
}

// 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)
	}
	return ip
}