require "util"
local bit = require "bit"

local p = Proto("ieee", "IEEE802.11 frame header")

-- based off <https://www.problemkaputt.de/gbatek.htm#dswifiieee80211frames>

p.fields.ctl = ProtoField.new("Frame control", "ieee.ctl", ftypes.BYTES)
p.fields.ctl_ver = ProtoField.uint16("ieee.ctl.ver", "Protocol version", base.DEC, {
	[0] = "Current",
	[1] = "Reserved",
	[2] = "Reserved",
	[3] = "Reserved",
}, bits(0, 2))
p.fields.ctl_type = ProtoField.uint16("ieee.ctl.type", "Type", base.DEC, {
	[0] = "Management",
	[1] = "Control",
	[2] = "Data",
	[3] = "Reserved",
}, bits(2, 2))
p.fields.ctl_subtype = ProtoField.uint16("ieee.ctl.subtype", "Subtype", base.DEC, nil, bits(4, 4))
p.fields.ctl_to_ds = ProtoField.bool("ieee.ctl.tods", "To DS", base.DEC, nil, bits(8))
p.fields.ctl_from_ds = ProtoField.bool("ieee.ctl.fromds", "From DS", base.DEC, nil, bits(9))
p.fields.ctl_fragment = ProtoField.bool("ieee.ctl.fragment", "More fragments", base.DEC, nil, bits(10))
p.fields.ctl_retry = ProtoField.bool("ieee.ctl.retry", "Retry", base.DEC, nil, bits(11))
p.fields.ctl_power = ProtoField.bool("ieee.ctl.power", "Power management", base.DEC, nil, bits(12))
p.fields.ctl_data = ProtoField.bool("ieee.ctl.data", "More data", base.DEC, nil, bits(13))
p.fields.ctl_wep = ProtoField.bool("ieee.ctl.wep", "WEP encrypt", base.DEC, nil, bits(14))
p.fields.ctl_order = ProtoField.bool("ieee.ctl.order", "Order", base.DEC, nil, bits(15))

p.fields.duration = ProtoField.uint16("ieee.duration", "Duration / ID")
p.fields.addr1 = ProtoField.ether("ieee.addr1", "Address 1")
p.fields.addr2 = ProtoField.ether("ieee.addr2", "Address 2")
p.fields.addr3 = ProtoField.ether("ieee.addr3", "Address 3")
p.fields.seq = ProtoField.uint16("ieee.seq", "Sequence control")
p.fields.seq_frag = ProtoField.uint16("ieee.seq.frag", "Fragment", base.DEC, nil, bits(0, 4))
p.fields.seq_num = ProtoField.uint16("ieee.seq.num", "Sequence number", base.DEC, nil, bits(4, 12))

p.fields.body = ProtoField.bytes("ieee.body", "Body")

local pc_dissector = Dissector.get("pictochat")

function p.dissector(buffer, pinfo, tree)
	local buffer_len = buffer:len()
	-- invalid
	if buffer_len < 10 then return 0 end

	-- pretty wireshark shit
	pinfo.cols.protocol = p.name

	-- MAC header is (usually) 0x18 bytes, but also sometimes contains values in
	-- the trailer. The 0x18 here is so wireshark only highlights the MAC header
	-- when clicking this item in the dissection tree.
	local subtree = tree:add(p, buffer(0x00, 0x18), p.description)
	local trailer_size = 0

	local ctl_tree = subtree:add_le(p.fields.ctl, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_ver, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_type, buffer(0x00, 2))
		local ctl_type = bit.rshift(bit.band(buffer(0x00, 2):le_uint(), bits(2, 2)), 2)
		ctl_tree:add_le(p.fields.ctl_subtype, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_to_ds, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_from_ds, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_fragment, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_retry, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_power, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_data, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_wep, buffer(0x00, 2))
		ctl_tree:add_le(p.fields.ctl_order, buffer(0x00, 2))

	subtree:add_le(p.fields.duration, buffer(0x02, 2))
	subtree:add(p.fields.addr1, buffer(0x04, 6))
	subtree:add(p.fields.addr2, buffer(0x0a, 6))
	subtree:add(p.fields.addr3, buffer(0x10, 6))
	local seq_tree = subtree:add_le(p.fields.seq, buffer(0x16, 2))
		seq_tree:add_le(p.fields.seq_frag, buffer(0x16, 2))
		seq_tree:add_le(p.fields.seq_num, buffer(0x16, 2))

	if ctl_type ~= 0 then
		trailer_size = 4 -- Frame Check Sequence (FCS) (hardware-generated CRC32)
	end

	buffer = buffer(0x18) -- seek forward

	local body_size = buffer:len() - trailer_size
	subtree:add(p.fields.body, buffer(0, body_size))
	pc_dissector:call(buffer(0, body_size):tvb(), pinfo, tree)

	if trailer_size == 0 then
		return buffer_len
	end

	buffer = buffer(body_size)

	return buffer_len
end