TUN/TAP packet decoder prototype

tuntap.py

import struct

import bitstream
import dotdict
import numpy as np
import pytun

def read_uint(stream, n=None):
	integer = 0
	for _ in range(n):
		integer <<= 1
		if stream.read(bool):
			integer += 1
	return integer
bitstream.register(int, read_uint)

def hexprint(buf):
	pbuf = buf
	while True:
		if len(pbuf) == 0:
			break

		linesize = min(len(pbuf), 8)
		for i in range(8):
			if i < linesize:
				print("%02x " % (pbuf[i],), end="")
			else:
				print("   ", end="")
		for i in range(linesize):
			print("%s" % (chr(pbuf[i]) if 32 <= pbuf[i] <= 127 else '.'), end="")
		print("")

		pbuf = pbuf[linesize:]

def decode_ether(bs):
	ether = dotdict.dotdict()

	ether.dst = ':'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, 6)))
	ether.src = ':'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, 6)))
	ether.type = bs.read(np.uint16)

	print("Ether %s -> %s Type: %04x" % (ether.src, ether.dst, ether.type))

	return ether

def decode_arp(bs):
	arp = dotdict.dotdict()

	arp.hrd = bs.read(np.uint16)
	arp.pro = bs.read(np.uint16)
	arp.hln = bs.read(np.uint8)
	arp.pln = bs.read(np.uint8)
	arp.op = bs.read(np.uint16)

	arp.sha = '.'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, arp.hln)))
	arp.spa = '.'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, arp.pln)))
	arp.tha = '.'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, arp.hln)))
	arp.tpa = '.'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, arp.pln)))

	print("ARP %04x/%04x %s %s -> %s (%s -> %s)" % (arp.hrd, arp.pro, "req" if arp.op == 1 else "rep" if arp.op == 2 else "???", arp.spa, arp.tpa, arp.sha, arp.tha))

	return arp

def decode_ip(bs):
	ip = dotdict.dotdict()

	ip.version = bs.read(int, 4)

	if ip.version == 4:
		return decode_ipv4(ip, bs)
	elif ip.version == 6:
		return decode_ipv6(ip, bs)
	else:
		return None

def decode_ipv4(ip, bs):
	ip.ihl = bs.read(int, 4) * 4
	ip.tos = str(bs.read(n=8))
	ip.length = bs.read(np.uint16)

	ip.id = bs.read(np.uint16)
	ip.flags = bs.read(int, 3)
	ip.frag = bs.read(int, 13) * 8

	ip.ttl = bs.read(np.uint8)
	ip.proto = bs.read(np.uint8)
	ip.checksum = bs.read(np.uint16)

	ip.src = '.'.join(bs.read(np.uint8, 4).astype(str))
	ip.dst = '.'.join(bs.read(np.uint8, 4).astype(str))

	ip.options_padding = bs.read(n=ip.ihl-20)

	ip.payload_length = ip.length-ip.ihl

	print("IPv%d %s -> %s Len: %d+%d Type: %d" % (ip.version, ip.src, ip.dst, ip.ihl, ip.length-ip.ihl, ip.proto))

	return ip

def decode_ipv6(ip, bs):
	ip.tclass = bs.read(np.uint8)
	ip.flow_label = bs.read(int, 20)

	ip.length = bs.read(np.uint16)
	ip.proto = bs.read(np.uint8)
	ip.ttl = bs.read(np.uint8)

	ip.src = ':'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, 16)))
	ip.dst = ':'.join(np.vectorize(lambda n: "%02x" % (n,))(bs.read(np.uint8, 16)))

	ip.payload_length = ip.length

	print("IPv%d %s -> %s Len: 40+%d Type: %d" % (ip.version, ip.src, ip.dst, ip.length, ip.proto))

	return ip

def decode_udp(bs):
	udp = dotdict.dotdict()

	udp.sport = bs.read(np.uint16)
	udp.dport = bs.read(np.uint16)
	udp.length = bs.read(np.uint16) - 8
	udp.checksum = bs.read(np.uint16)
	udp.data = bs.read(bytes, udp.length)

	print("UDP %d -> %d Len: 8+%d" % (udp.sport, udp.dport, udp.length))

	return udp

def main():
	devname = 'test2'
	tap = True

	dev = pytun.TunTapDevice(devname, flags=pytun.IFF_TAP if tap else pytun.IFF_TUN)

	while True:
		buf = dev.read(4096)

		print("---")

		bs = bitstream.BitStream(buf)
		tuntap_flags = bs.read(np.uint16)
		tuntap_proto = bs.read(np.uint16)
		print("%s Flags: %04x Proto: %04x Len: 4+%d" % ("TAP" if tap else "TUN", tuntap_flags, tuntap_proto, len(buf)-4))

		print("-")
		hexprint(buf[4:])
		print("-")

		data = None

		if tap:
			ether = decode_ether(bs)

			if ether.type == 0x0806:
				arp = decode_arp(bs)

		if not tap or ether.type == 0x0800 or ether.type == 0x86dd:
			ip = decode_ip(bs)

			if ip.proto == 17:
				udp = decode_udp(bs)
				data = udp.data

		if data is not None:
			print("-")
			hexprint(data)

		print("---")
		print("")

if __name__ == '__main__':
	main()