// TODO: slice bounds checking
// TODO: authentication
// TODO: somehow persist IPs
// TODO: handle multiple peers behind one NAT
package main

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

	"baalvpn/shared"
)

type Peer struct {
	RealAddr   *net.UDPAddr
	InternalIP string
}

// TODO: replace with key exchange between peers
const PSK = "rj5FvWB0n342C6UGrAS0gI3GX5ya+lDmC7/XmgFHF+s="

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

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

	var err error
	pskKey, err = base64.StdEncoding.DecodeString(PSK)
	if err != nil {
		panic(err)
	}

	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()

		internalIP := randomIP().String()

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

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

		pkt, err := shared.DecryptSym(pskKey, req[4:])
		if err != nil {
			log.Printf("failed to decrypt packet from %s: %s", peer.RealAddr, err)
			return
		}

		etherType := binary.BigEndian.Uint16(pkt[12:14])
		payload := pkt[14:]

		switch etherType {
		case 0x86dd: // IPv6
			rawDestIP := payload[24:40]
			destIP := net.IP(rawDestIP).String()

			encryptedPkt, err := shared.EncryptSym(pskKey, pkt)
			if err != nil {
				panic(err)
			}
			resp := shared.BuildPkt(shared.RESP_DATA, encryptedPkt)

			if rawDestIP[0] == 0xff { // multicast
				log.Println(peer.InternalIP + " -> " + destIP + " (*)")
				broadcast(resp, peer.InternalIP)
			} else if destPeer := getPeerByInternalIP(destIP); destPeer != nil {
				log.Println(peer.InternalIP + " -> " + destIP)
				sendTo(destPeer.RealAddr, resp)
			} else {
				log.Println(peer.InternalIP + " -/> " + destIP + " (unrecognized IP)")
			}
		default:
			log.Printf("unhandled EtherType: 0x%x\n", etherType)
		}
	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
}