/*
 * Insert hops in traceroute (ipv6 only)
 *
 * Create a virtual interface, and send "time exceeded" messages
 * until the TTL match a specified number, effectively adding hops
 * to a traceroute.
 *
 * The source of the error messages is changed so the traceroute
 * appears to be incrementing, up to the requested destination.
 *
 * by wgs
 */

package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"log"
	"net"

	"golang.org/x/net/icmp"
	"golang.zx2c4.com/wireguard/tun"
)

const (
	TUN_HEADER  = 4
	IPV6_HEADER = 40
	ICMP_HEADER = 4
)

/*
 * Compute IPv6 checksum for a given packet
 * https://datatracker.ietf.org/doc/html/rfc4443#section-2.3
 */
func checksum(body []byte, srcIP, dstIP net.IP) (crc []byte) {
	out := make([]byte, 2)
	// from golang.org/x/net/icmp/message.go
	checksum := func(b []byte) uint16 {
		csumcv := len(b) - 1 // checksum coverage
		s := uint32(0)
		for i := 0; i < csumcv; i += 2 {
			s += uint32(b[i+1])<<8 | uint32(b[i])
		}
		if csumcv&1 == 0 {
			s += uint32(b[csumcv])
		}
		s = s>>16 + s&0xffff
		s = s + s>>16
		return ^uint16(s)
	}

	b := body
	l := len(b)
	psh := icmp.IPv6PseudoHeader(srcIP, dstIP)
	b = append(psh, b...)
	off := 2 * net.IPv6len
	binary.BigEndian.PutUint32(b[off:off+4], uint32(l))
	s := checksum(b)
	out[0] ^= byte(s)
	out[1] ^= byte(s >> 8)

	return out
}

func ipv6_header(src, dst net.IP, len uint16) []byte {
	header := make([]byte, IPV6_HEADER)

	// Packet version (ipv6), Traffic class and Flow label
	copy(header[:4], []byte{0x60, 0, 0, 0})

	sz := new(bytes.Buffer)
	binary.Write(sz, binary.BigEndian, len)
	header[4] = sz.Bytes()[0] // Packet size
	header[5] = sz.Bytes()[1] // Packet size
	header[6] = 0x3a          // Next header (58, ICMPv6)
	header[7] = 0x40          // Hop Limit (64)

	// source / destination ipv6
	copy(header[8:], src[:16])
	copy(header[24:], dst[:16])

	return header
}

func icmp_message(icmp_type, icmp_code uint8, src, dst net.IP, payload []byte) []byte {
	message := make([]byte, ICMP_HEADER)

	message[0] = icmp_type
	message[1] = icmp_code

	switch icmp_type {
	case 1: // Destination unreachable
		fallthrough
	case 3: // Time Exceeded
		message = append(message, 0, 0, 0, 0) // unused
		message = append(message, payload...)
	case 129: // Echo reply
		// payload must include identifier + seqnum
		// from original message
		message = append(message, payload...)
	}

	crc := checksum(message, src, dst)
	message[2] = crc[0]
	message[3] = crc[1]

	return message
}

func main() {
	ifname := flag.String("i", "tun", "Interface name")
	mtu := flag.Int("m", 1500, "Interface MTU")
	hop := flag.Int("n", 4, "Hops number to insert")
	verbose := flag.Bool("v", false, "Enable verbose logging")

	flag.Parse()

	tun, err := tun.CreateTUN(*ifname, *mtu)
	if err != nil {
		log.Fatal(err)
	}

	hops := uint8(*hop)

	for {
		var icmp, ipv6 []byte

		src := make([]byte, 16)
		dst := make([]byte, 16)
		buf := make([]byte, *mtu+TUN_HEADER)
		sz, err := tun.Read(buf, TUN_HEADER)
		if err != nil {
			log.Fatalf("Read %s: %s", *ifname, err.Error())
		}

		// Invalid packet, ignore it
		if sz < IPV6_HEADER {
			continue
		}

		packet := buf[TUN_HEADER : TUN_HEADER+sz]

		/*
		 * Skip packet if the specified number of hops cannot
		 * be inserted by only changing the last byte
		 */
		if packet[39] < hops {
			if *verbose == true {
				log.Printf("%s %s > %s Dropped (last byte 0x%02x < 0x%02x)",
					*ifname,
					net.IP(packet[24:40]).String(),
					net.IP(packet[8:24]).String(),
					packet[39], hops-1)
			}
			continue
		}

		ttl := uint8(packet[7])

		// Invert source and destination address
		copy(dst, packet[8:8+16])
		copy(src, packet[24:24+16])

		/*
		 * If TTL is lower than the configured number of hops,
		 * start sending ICMP time exceeded replies to the
		 * originating source.
		 */
		if ttl < hops+1 {
			/*
			 * Use hexa representation of TTL as the
			 * last byte of source address, thus
			 * incrementing hops until final destination
			 * is reached
			 */
			src[15] = src[15] - hops + ttl

			// ICMP error must fit in 1280 bytes (ipv6 min. mtu)
			if sz+IPV6_HEADER+ICMP_HEADER > 1280 {
				sz = 1228
			}

			ipv6 = ipv6_header(src, dst, uint16(sz+ICMP_HEADER+4))

			if ttl < hops {
				// Time exceeded
				icmp = icmp_message(3, 0, src, dst, packet[:sz])
			} else {
				// Destination port unreachable
				icmp = icmp_message(1, 4, src, dst, packet[:sz])
			}
		} else {
			// Skip everything not ICMPv6
			if packet[6] != 58 {
				continue
			}

			switch packet[IPV6_HEADER] {
			case 128: // ICMPv6 echo request
				ipv6 = ipv6_header(src, dst, uint16(sz-IPV6_HEADER))
				icmp = icmp_message(129, 0, src, dst, packet[IPV6_HEADER+4:])
			default:
				continue
			}
		}

		// Construct final return packet
		b := make([]byte, TUN_HEADER)
		b = append(b, ipv6...)
		b = append(b, icmp...)

		if *verbose == true {
			log.Printf("%s %s > %s ICMP Type 0x%02x Code %d (%d bytes)",
				*ifname,
				net.IP(src).String(),
				net.IP(dst).String(),
				icmp[0],
				icmp[1],
				len(b))
		}

		tun.Write(b, TUN_HEADER)
	}
}