wacky encryption

client/main.go

@@ -1,23 +1,24 @@
 package main
 
 import (
-	"encoding/base64"
+	"crypto/sha256"
 	"encoding/binary"
 	"encoding/hex"
 	"log"
 	"net"
 	"os"
 	"os/exec"
+	"time"
 
 	"baalvpn/shared"
 
+	"filippo.io/mlkem768/xwing"
 	"github.com/songgao/water"
 )
 
 const (
 	SERVER_IP  = "172.20.12.47"
 	IFACE_NAME = "baalvpn"
-	PSK        = "rj5FvWB0n342C6UGrAS0gI3GX5ya+lDmC7/XmgFHF+s="
 )
 
 var (
@@ -25,7 +26,14 @@ var (
 	serverAddr *net.UDPAddr
 	iface      *water.Interface
 	conn       *net.UDPConn
-	pskKey     []byte
+	privKey    *xwing.DecapsulationKey
+
+	// TODO: this should probably have a mutex
+	peerPubKeys = map[string][]byte{}
+	// TODO: this should probably have a mutex
+	peerSessionKeys = map[string][]byte{}
+	// TODO: this should probably have a mutex
+	peerEstablishAcks = map[string]bool{}
 )
 
 func main() {
@@ -33,12 +41,6 @@ func main() {
 		panic("root permissions needed")
 	}
 
-	var err error
-	pskKey, err = base64.StdEncoding.DecodeString(PSK)
-	if err != nil {
-		panic(err)
-	}
-
 	listenAddr, err := net.ResolveUDPAddr("udp", ":38000")
 	if err != nil {
 		panic(err)
@@ -61,12 +63,63 @@ func main() {
 }
 
 func register() {
-	req := shared.BuildPkt(shared.REQ_REGISTER, []byte{})
+	var err error
+	privKey, err = xwing.GenerateKey()
+	if err != nil {
+		panic(err)
+	}
+	pubKey := privKey.EncapsulationKey()
+
+	req := shared.BuildPkt(shared.REQ_REGISTER, pubKey)
 	if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
 		panic(err)
 	}
 }
 
+func getOrEstablishSessionKey(ip string) []byte {
+	if key, ok := peerSessionKeys[ip]; ok {
+		return key
+	} else {
+		log.Println("requesting pubkey of " + ip)
+		req := shared.BuildPkt(shared.REQ_GET_PUBKEY, net.ParseIP(ip))
+		if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+			panic(err)
+		}
+
+		for {
+			// TODO: eww
+			time.Sleep(50 * time.Millisecond)
+			if peerPubKey, ok := peerPubKeys[ip]; ok {
+				log.Println("got pubkey of " + ip)
+
+				ciphertext, sharedSecret, err := xwing.Encapsulate(peerPubKey)
+				if err != nil {
+					panic(err)
+				}
+
+				reqData := append(net.ParseIP(ip), append(internalIP, ciphertext...)...)
+				req := shared.BuildPkt(shared.REQ_ESTABLISH, reqData)
+				if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+					panic(err)
+				}
+
+				for {
+					// TODO: eww
+					time.Sleep(50 * time.Millisecond)
+					if peerEstablishAcks[ip] {
+						delete(peerEstablishAcks, ip)
+						break
+					}
+				}
+
+				sessionKey := sha256.Sum256(sharedSecret)
+				peerSessionKeys[ip] = sessionKey[:]
+				return sessionKey[:]
+			}
+		}
+	}
+}
+
 func receivePackets() {
 	for {
 		buffer := make([]byte, 50000)
@@ -106,8 +159,12 @@ func receivePacket(pkt []byte) {
 		log.Println("got assigned IP by the server:", internalIP.String())
 		setupInterface()
 		go sendPackets()
-	case shared.RESP_DATA:
-		decryptedPkt, err := shared.DecryptSym(pskKey, pkt[4:])
+	case shared.ENC_PKT:
+		srcIP := net.IP(pkt[20:36]).String()
+
+		sessionKey := getOrEstablishSessionKey(srcIP)
+
+		decryptedPkt, err := DecryptSym(sessionKey, pkt[36:])
 		if err != nil {
 			panic(err)
 		}
@@ -116,6 +173,37 @@ func receivePacket(pkt []byte) {
 		if _, err := iface.Write(decryptedPkt); err != nil {
 			log.Println(err)
 		}
+	case shared.BROADCAST_PKT:
+		log.Println("received pkt:", hex.EncodeToString(pkt[4:]))
+		if _, err := iface.Write(pkt[4:]); err != nil {
+			log.Println(err)
+		}
+	case shared.RESP_GET_PUBKEY:
+		respIP := net.IP(pkt[4:20]).String()
+		pubKey := pkt[20:]
+		peerPubKeys[respIP] = pubKey
+	case shared.REQ_ESTABLISH:
+		srcIP := net.IP(pkt[20:36]).String()
+		ciphertext := pkt[36:]
+
+		sharedSecret, err := xwing.Decapsulate(privKey, ciphertext)
+		if err != nil {
+			panic(err)
+		}
+
+		sessionKey := sha256.Sum256(sharedSecret)
+		peerSessionKeys[srcIP] = sessionKey[:]
+
+		log.Println("received session key from", srcIP)
+		reqData := append(net.ParseIP(srcIP), internalIP...)
+		req := shared.BuildPkt(shared.RESP_ESTABLISH, reqData)
+		if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+			panic(err)
+		}
+	case shared.RESP_ESTABLISH:
+		srcIP := net.IP(pkt[20:36]).String()
+		log.Println("established session key with", srcIP)
+		peerEstablishAcks[srcIP] = true
 	default:
 		log.Println("unknown packet type")
 	}
@@ -150,14 +238,33 @@ func sendPackets() {
 		}
 		pkt = pkt[:n]
 
-		encryptedPkt, err := shared.EncryptSym(pskKey, pkt)
-		if err != nil {
-			panic(err)
+		etherType := binary.BigEndian.Uint16(pkt[12:14])
+		payload := pkt[14:]
+
+		if etherType != 0x86dd { // IPv6
+			continue
 		}
 
-		req := shared.BuildPkt(shared.REQ_DATA, encryptedPkt)
-		if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
-			log.Println(err)
+		destIP := net.IP(payload[24:40]).String()
+
+		if payload[24] == 0xff { // multicast
+			req := shared.BuildPkt(shared.BROADCAST_PKT, pkt)
+			if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+				log.Println(err)
+			}
+		} else {
+			destSessionKey := getOrEstablishSessionKey(destIP)
+
+			encryptedPkt, err := EncryptSym(destSessionKey, pkt)
+			if err != nil {
+				panic(err)
+			}
+
+			reqBody := append(net.ParseIP(destIP), append(internalIP, encryptedPkt...)...)
+			req := shared.BuildPkt(shared.ENC_PKT, reqBody)
+			if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+				log.Println(err)
+			}
 		}
 	}
 }