Quick, very dirty, wg ping demo

wg_demo.go

/* Copyright 2017 rot256
 *
 * Permission is hereby granted, free of charge,
 * to any person obtaining a copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"github.com/aead/blake2s"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"hash"
	"log"
	"net"
	"os"
	"time"
)

const (
	TagSize       = 16 // Poly1305
	MACSize       = 16
	HashSize      = 32 // Blake2s 256-bit output
	AEADKeySize   = 32
	AEADNonceSize = 12 // Rest is zero
	AEADCntSize   = 8
	TimestampSize = 12
	PublicKeySize = 32 // curve25519 public key size (point)
	SecretKeySize = 32 // curve25519 secret key size (scalar)
)

type MessageInital struct {
	Type      uint32
	Sender    uint32
	Ephemeral [PublicKeySize]byte
	Static    [PublicKeySize + TagSize]byte
	Timestamp [TimestampSize + TagSize]byte
	// followed by macs
}

type MessageResponse struct {
	Type      uint32
	Sender    uint32
	Reciever  uint32
	Ephemeral [PublicKeySize]byte
	Empty     [TagSize]byte
	Mac1      [MACSize]byte
	Mac2      [MACSize]byte
}

type MessageTransport struct {
	Type     uint32
	Reciever uint32
	Counter  uint64
}

type Macs struct {
	Mac1 [MACSize]byte
	Mac2 [MACSize]byte
}

const (
	Construction = "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
	Identifier   = "WireGuard v1 zx2c4 Jason@zx2c4.com"
	LabelMAC1    = "mac1----"
	LabelCookie  = "cookie--"
)

var RemoteStaticPK [PublicKeySize]byte
var LocalStaticPK [PublicKeySize]byte
var LocalStaticSK [SecretKeySize]byte
var Preshared [HashSize]byte

func ipChecksum(buf []byte) uint16 {
	sum := uint32(0)
	for ; len(buf) >= 2; buf = buf[2:] {
		sum += uint32(buf[0])<<8 | uint32(buf[1])
	}
	if len(buf) > 0 {
		sum += uint32(buf[0]) << 8
	}
	for sum > 0xffff {
		sum = (sum >> 16) + (sum & 0xffff)
	}
	csum := ^uint16(sum)
	if csum == 0 {
		csum = 0xffff
	}
	return csum
}

func init() {
	// demo.wireguard
	ourPrivate, _ := base64.StdEncoding.DecodeString("WAmgVYXkbT2bCtdcDwolI88/iVi/aV3/PHcUBTQSYmo=")
	ourPublic, _ := base64.StdEncoding.DecodeString("K5sF9yESrSBsOXPd6TcpKNgqoy1Ik3ZFKl4FolzrRyI=")
	theirPublic, _ := base64.StdEncoding.DecodeString("qRCwZSKInrMAq5sepfCdaCsRJaoLe5jhtzfiw7CjbwM=")
	preshared, _ := base64.StdEncoding.DecodeString("FpCyhws9cxwWoV4xELtfJvjJN+zQVRPISllRWgeopVE=")

	copy(RemoteStaticPK[:], theirPublic)
	copy(LocalStaticSK[:], ourPrivate)
	copy(LocalStaticPK[:], ourPublic)
	copy(Preshared[:], preshared)

	var temp [PublicKeySize]byte
	curve25519.ScalarBaseMult(&temp, &LocalStaticSK)
	if temp != LocalStaticPK {
		panic(nil)
	}
	fmt.Println("Loaded")
}

func Hash(val []byte) [HashSize]byte {
	return blake2s.Sum256(val)
}

func HMAC(key []byte, input []byte) [HashSize]byte {
	mac := hmac.New(func() hash.Hash {
		h, _ := blake2s.New256(nil)
		return h
	}, key)
	mac.Write(input)
	tmp := mac.Sum(nil)
	var sum [HashSize]byte
	copy(sum[:], tmp)
	return sum
}

func MAC(key []byte, input []byte) [MACSize]byte {
	var sum [MACSize]byte
	hsh, err := blake2s.New128(key)
	if err != nil {
		panic(err)
	}
	hsh.Write(input)
	tmp := hsh.Sum(nil)
	copy(sum[:], tmp)
	return sum
}

/*
 * ref: https://tools.ietf.org/html/rfc5869
 */

func KDF1(key []byte, input []byte) [HashSize]byte {
	prk := HMAC(key, input)
	t1 := HMAC(prk[:], []byte{0x1})
	return t1
}

func KDF2(key []byte, input []byte) (t1 [HashSize]byte, t2 [HashSize]byte) {
	prk := HMAC(key, input)
	t1 = HMAC(prk[:], []byte{0x1})
	t2 = HMAC(prk[:], append(t1[:], 0x2))
	return
}

func KDF3(key []byte, input []byte) (t1 [HashSize]byte, t2 [HashSize]byte, t3 [HashSize]byte) {
	prk := HMAC(key, input)
	t1 = HMAC(prk[:], []byte{0x1})
	t2 = HMAC(prk[:], append(t1[:], 0x2))
	t3 = HMAC(prk[:], append(t2[:], 0x3))
	return
}

func init() {
	key, _ := hex.DecodeString("60e26daef327efc02ec335e2a025d2d016eb4206f87277f52d38d1988b78cd36")
	input, _ := hex.DecodeString("e6f2a4d1c28ee5c7ad0329268255a468ad407d2672824c0c0eb30ea6ef450145")

	tmp, _ := hex.DecodeString("56b95e3748a3adb6c401615767be23622aee034e4414e8f9fbec58b342882fb4")
	var output [32]byte
	copy(output[:], tmp)

	out := KDF1(key, input)
	fmt.Println(hex.EncodeToString(out[:]))
	fmt.Println(hex.EncodeToString(tmp))
	if output != out {
		panic(nil)
	}

}

func Timestamp() [TimestampSize]byte {
	var tai64n [TimestampSize]byte
	now := time.Now()
	binary.BigEndian.PutUint64(tai64n[:], 4611686018427387914+uint64(now.Unix()))
	binary.BigEndian.PutUint32(tai64n[8:], uint32(now.UnixNano()))
	return tai64n
}

func AEADSeal(
	key [AEADKeySize]byte,
	cnt uint64,
	pt []byte,
	ad []byte,
) []byte {
	aead, err := chacha20poly1305.New(key[:])
	if err != nil {
		panic(err)
	}
	nonce := bytes.NewBuffer([]byte{0, 0, 0, 0})
	binary.Write(nonce, binary.LittleEndian, cnt)
	return aead.Seal([]byte{}, nonce.Bytes(), pt, ad)
}

func AEADOpen(
	key [AEADKeySize]byte,
	cnt uint64,
	ct []byte,
	ad []byte,
) ([]byte, error) {
	aead, err := chacha20poly1305.New(key[:])
	if err != nil {
		panic(err)
	}
	nonce := bytes.NewBuffer([]byte{0, 0, 0, 0})
	binary.Write(nonce, binary.LittleEndian, cnt)
	return aead.Open([]byte{}, nonce.Bytes(), ct, ad)
}

func makeMessageInitial() ([]byte, [HashSize]byte, [HashSize]byte, [PublicKeySize]byte) {

	var C [HashSize]byte
	var H [HashSize]byte
	var k [HashSize]byte

	C = Hash([]byte(Construction))
	H = Hash(append(C[:], []byte(Identifier)...))
	H = Hash(append(H[:], RemoteStaticPK[:]...))

	var eSK [SecretKeySize]byte
	var ePK [PublicKeySize]byte

	// generate ephemeral keypair
	rand.Read(eSK[:])
	curve25519.ScalarBaseMult(&ePK, &eSK)

	C = KDF1(C[:], ePK[:])

	var msg MessageInital

	msg.Type = 1
	msg.Sender = 28
	msg.Ephemeral = ePK

	H = Hash(append(H[:], ePK[:]...))

	func() {
		var ss [PublicKeySize]byte
		curve25519.ScalarMult(&ss, &eSK, &RemoteStaticPK)
		C, k = KDF2(C[:], ss[:])
	}()

	func() {
		tmp := AEADSeal(k, 0, LocalStaticPK[:], H[:])
		copy(msg.Static[:], tmp)
	}()

	H = Hash(append(H[:], msg.Static[:]...))

	func() {
		var ss [32]byte
		curve25519.ScalarMult(&ss, &LocalStaticSK, &RemoteStaticPK)
		C, k = KDF2(C[:], ss[:])
	}()

	func() {
		tmp1 := Timestamp()
		tmp2 := AEADSeal(k, 0, tmp1[:], H[:])
		copy(msg.Timestamp[:], tmp2)
	}()
	H = Hash(append(H[:], msg.Timestamp[:]...))

	packet := make([]byte, 0, 256)
	writer := bytes.NewBuffer(packet)
	binary.Write(writer, binary.LittleEndian, &msg)

	func() {
		var macs Macs
		key := Hash(append([]byte(LabelMAC1), RemoteStaticPK[:]...))
		macs.Mac1 = MAC(key[:], writer.Bytes())
		binary.Write(writer, binary.LittleEndian, &macs)
	}()

	return writer.Bytes(), C, H, eSK
}

func processMessageResponse(resp []byte, C [HashSize]byte, H [HashSize]byte, eSK [32]byte) (keySend [AEADKeySize]byte, keyRecv [AEADKeySize]byte, recv uint32) {

	// process response

	var msgResp MessageResponse
	reader := bytes.NewReader(resp)
	binary.Read(reader, binary.LittleEndian, &msgResp)

	if msgResp.Type != 2 {
		panic(nil)
	}

	C = KDF1(C[:], msgResp.Ephemeral[:])
	H = Hash(append(H[:], msgResp.Ephemeral[:]...))

	func() {
		var ss [32]byte
		curve25519.ScalarMult(&ss, &eSK, &msgResp.Ephemeral)
		C = KDF1(C[:], ss[:])
	}()

	func() {
		var ss [32]byte
		curve25519.ScalarMult(&ss, &LocalStaticSK, &msgResp.Ephemeral)
		C = KDF1(C[:], ss[:])
	}()

	var tau [HashSize]byte
	var k [HashSize]byte

	C, tau, k = KDF3(C[:], Preshared[:])

	H = Hash(append(H[:], tau[:]...))

	_, err := AEADOpen(k, 0, msgResp.Empty[:], H[:])
	if err != nil {
		panic(err)
	}

	H = Hash(append(H[:], msgResp.Empty[:]...))

	// derive final symmetric keys
	if msgResp.Reciever != 28 {
		panic(nil)
	}
	recv = msgResp.Sender
	keySend, keyRecv = KDF2(C[:], []byte{})
	return
}

// padding ommitted for simplicity
func makeMessageTransport(data []byte, key [AEADKeySize]byte, recv uint32, cnt uint64) []byte {
	var msg MessageTransport

	msg.Type = 4
	msg.Reciever = recv
	msg.Counter = cnt
	payload := AEADSeal(key, cnt, data, []byte{})

	packet := make([]byte, 0, 256)
	writer := bytes.NewBuffer(packet)
	binary.Write(writer, binary.LittleEndian, msg)
	writer.Write(payload)

	return writer.Bytes()
}

func processMessageTransport(resp []byte, key [AEADKeySize]byte) []byte {
	var msg MessageTransport
	reader := bytes.NewReader(resp)
	binary.Read(reader, binary.LittleEndian, &msg)
	payload, err := AEADOpen(key, msg.Counter, resp[binary.Size(msg):], []byte{})
	if err != nil {
		panic(err)
	}
	return payload
}

func test() error {
	dst := os.Args[1]
	fmt.Println(dst)
	serverAddr, err := net.ResolveUDPAddr("udp", dst)
	if err != nil {
		return err
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}
	defer conn.Close()

	// send frist message
	msg, C, H, eSK := makeMessageInitial()
	n, err := conn.Write(msg)
	if err != nil {
		return err
	}

	// read response
	buf := make([]byte, 1024)
	n, addr, err := conn.ReadFromUDP(buf)
	fmt.Println("Received ", hex.EncodeToString(buf[0:n]), " from ", addr)
	if err != nil {
		return err
	}
	keySend, keyRecv, recv := processMessageResponse(buf[0:n], C, H, eSK)

	fmt.Println(keySend)
	fmt.Println(keyRecv)
	fmt.Println(recv)

	// write ICMP Echo packet
	pingMessage, _ := (&icmp.Message{
		Type: ipv4.ICMPTypeEcho,
		Body: &icmp.Echo{
			ID:   921,
			Seq:  438,
			Data: []byte("WireGuard"),
		},
	}).Marshal(nil)
	pingHeader, err := (&ipv4.Header{
		Version:  ipv4.Version,
		Len:      ipv4.HeaderLen,
		TotalLen: ipv4.HeaderLen + len(pingMessage),
		Protocol: 1, // ICMP
		TTL:      20,
		Src:      net.IPv4(10, 189, 129, 2),
		Dst:      net.IPv4(10, 189, 129, 1),
	}).Marshal()
	binary.BigEndian.PutUint16(pingHeader[2:], uint16(ipv4.HeaderLen+len(pingMessage))) // fix the length endianness on BSDs
	pingData := append(pingHeader, pingMessage...)
	binary.BigEndian.PutUint16(pingData[10:], ipChecksum(pingData))

	// create transport message
	msg = makeMessageTransport(pingData, keySend, recv, 0)
	fmt.Println("Transport message:", hex.EncodeToString(msg))
	n, err = conn.Write(msg)
	if err != nil {
		return err
	}

	// read response
	n, addr, err = conn.ReadFromUDP(buf)
	fmt.Println("Received ", hex.EncodeToString(buf[0:n]), " from ", addr)
	if err != nil {
		return err
	}

	// I stole this
	replyPacket := buf[:n]
	if replyPacket[0] != 4 { // Type: Data
		log.Fatalf("unexpected reply packet type: %d", replyPacket[0])
	}
	if replyPacket[1] != 0 || replyPacket[2] != 0 || replyPacket[3] != 0 {
		log.Fatalf("reply packet has non-zero reserved fields")
	}
	replyPacket = processMessageTransport(replyPacket, keyRecv)

	replyHeaderLen := int(replyPacket[0]&0x0f) << 2
	replyLen := binary.BigEndian.Uint16(replyPacket[2:])
	replyMessage, err := icmp.ParseMessage(1, replyPacket[replyHeaderLen:replyLen])
	if err != nil {
		log.Fatalf("error parsing echo: %s", err)
	}
	echo, ok := replyMessage.Body.(*icmp.Echo)
	if !ok {
		log.Fatalf("unexpected reply body type %T", replyMessage.Body)
	}

	if echo.ID != 921 || echo.Seq != 438 || string(echo.Data) != "WireGuard" {
		log.Fatalf("incorrect echo response: %#v", echo)
	}

	log.Println(echo)

	return nil
}

func main() {
	err := test()
	fmt.Println(err)
}

hello, what license of this code?

Added MIT license.
If you are interested in WireGuard for userspace you might want to take a look at my wireguard-go project: https://git.zx2c4.com/wireguard-go/about/