implement non-janky message reassembly

1 files changed, 105 insertions, 87 deletions

diff --git a/wireshark/pictochat.lua b/wireshark/pictochat.lua
index ff6c21f..655f27e 100644
--- a/wireshark/pictochat.lua
+++ b/wireshark/pictochat.lua
@@ -8,166 +8,184 @@ function p.init()
 end
 
 local MSG_TYPE = {
-	INIT = 0,
-	MSG_START = 10,
-	MSG_END = 24,
-	USER_JOIN = 48,
-	STATUS = 52,
-	MSG_BODY = 86,
+	MSG_START = 1,
+	MSG_BODY = 2,
+	STATUS_END = 4,
+	STATUS = 5,
 }
 
 local MSG_TYPE_MAP = {
-	[MSG_TYPE.INIT] = "Init", -- sent once (init?)
-	[MSG_TYPE.MSG_START] = "Message start", -- sent before message body and user join
-	[MSG_TYPE.MSG_END] = "Message end", -- (sometimes) sent after multi-frame message body packets
-	[MSG_TYPE.USER_JOIN] = "User join", -- contain nickname / message
+	[MSG_TYPE.MSG_START] = "Message start",
+	[MSG_TYPE.MSG_BODY] = "Message body",
+	[MSG_TYPE.STATUS_END] = "Status (end)",
 	[MSG_TYPE.STATUS] = "Status", -- constantly spammed, contains room users' addresses
-	[MSG_TYPE.MSG_BODY] = "Message body", -- contains tile data for drawing
 }
 
 p.fields.unknown = ProtoField.uint16("pictochat.unknown", "Unknown")
 
-p.fields.payload_len = ProtoField.uint16("pictochat.payload_len", "Payload length")
-local function dissect_payload(buffer, pinfo, tree)
-	tree:add_le(p.fields.payload_len, buffer(0x00, 2))
-	local payload_length = buffer(0x00, 2):le_uint()
-	buffer = buffer(0x02)
-	return buffer, tree:add(buffer(0x00, payload_length), "Payload: " .. payload_length .. " bytes")
-end
-
 local dissect_msg_type = { }
 
+local state = {
+	pid_mid_map = {}, -- map packet id to pictochat message id (key = pid)
+	msg = {}, -- pictochat messages (key = mid)
+	mid = 0, -- current pictochat message id
+}
+
 p.fields.msg_start_len = ProtoField.uint8("pictochat.msg_start.len", "Total length")
-local state = {}
+p.fields.msg_start_id = ProtoField.uint32("pictochat.msg_start.id", "ID")
+p.fields.msg_start_seq = ProtoField.uint16("pictochat.msg_start.seq", "Sequence")
+p.fields.msg_start_new = ProtoField.bool("pictochat.msg_start.new", "New")
 dissect_msg_type[MSG_TYPE.MSG_START] = function (buffer, pinfo, tree)
-	local payload
-	buffer, payload = dissect_payload(buffer, pinfo, tree)
-
-	payload:add_le(p.fields.msg_start_len, buffer(0x04, 2))
-	local start_len = buffer(0x04, 2):le_uint()
+	local subtree = tree:add(p, buffer(), string.format("PictoChat Message start: %d bytes", buffer:len()))
+
+	subtree:add_le(p.fields.unknown, buffer(0x00, 2))
+	subtree:add_le(p.fields.unknown, buffer(0x02, 2)) -- usu. ffffh
+	subtree:add_le(p.fields.msg_start_len, buffer(0x04, 2))
+	local len = buffer(0x04, 2):le_uint()
+	subtree:add_le(p.fields.unknown, buffer(0x06, 2)) -- usu. 0000h
+	subtree:add_le(p.fields.unknown, buffer(0x08, 2))
+	subtree:add_le(p.fields.unknown, buffer(0x0a, 2)) -- usu. 0000h
+	subtree:add_le(p.fields.msg_start_id, buffer(0x0c, 4))
+	local id = buffer(0x0c, 4):le_uint()
+	subtree:add_le(p.fields.msg_start_seq, buffer(0x10, 2))
+	subtree:add_le(p.fields.new, buffer(0x12, 2))
+	local new = buffer(0x12, 2):le_uint() ~= 0
+
+	if new then
+		state.mid = id
+
+		if not pinfo.visited then
+			state.msg[state.mid] = {
+				buf = ByteArray.new(),
+				size = len,
+			}
+		end
+	end
 
-	state.body_remaining = start_len
-	state.buf = ByteArray.new()
-	state.want_next = 0
-end
+	state.pid_mid_map[pinfo.number] = state.mid
 
-p.fields.user_addr = ProtoField.bytes("pictochat.user.addr", "Address", base.COLON)
-p.fields.user_name = ProtoField.string("pictochat.user.name", "Nickname")
-p.fields.user_msg = ProtoField.string("pictochat.user.msg", "Message")
-p.fields.user_color = ProtoField.uint16("pictochat.user.color", "Color", base.DEC, PROFILE_COLOR_MAP)
-p.fields.user_bday_month = ProtoField.uint8("pictochat.user.bday_month", "Month")
-p.fields.user_bday_day = ProtoField.uint8("pictochat.user.bday_day", "Day")
-dissect_msg_type[MSG_TYPE.USER_JOIN] = function (buffer, pinfo, tree)
-	local payload
-	buffer, payload = dissect_payload(buffer, pinfo, tree)
-
-	add_addr_le(payload, p.fields.user_addr, buffer(0x0a, 6))
-	local user_name = buffer(0x10, 20):le_ustring()
-	payload:add(p.fields.user_name, buffer(0x10, 20), user_name)
-	register_addr_le(buffer(0x0a, 6):raw(), user_name)
-	pinfo.cols.info = string.format("%s (%s)", pinfo.cols.info, user_name)
-	payload:add(p.fields.user_msg, buffer(0x24, 52), buffer(0x24, 52):le_ustring())
-	payload:add_le(p.fields.user_color, buffer(0x58, 2))
-	local bday = payload:add(buffer(0x5a, 2), string.format("Birthday: %02d/%02d", buffer(0x5a, 1):uint(), buffer(0x5b, 1):uint()))
-	bday:add(p.fields.user_bday_month, buffer(0x5a, 1))
-	bday:add(p.fields.user_bday_day, buffer(0x5b, 1))
+	if pinfo.visited then
+		ByteArray.tvb(state.msg[state.pid_mid_map[pinfo.number]].buf, "Complete message")
+	end
 end
 
 p.fields.body_len = ProtoField.uint8("pictochat.body.len", "Length")
 p.fields.body_end = ProtoField.bool("pictochat.body.end", "End")
 p.fields.body_offset = ProtoField.uint16("pictochat.body.offset", "Offset")
-p.fields.body_seq = ProtoField.uint16("pictochat.body.seq", "Sequence")
 p.fields.body = ProtoField.bytes("pictochat.body", "Data")
+p.fields.body_seq = ProtoField.uint16("pictochat.body.seq", "Sequence")
+p.fields.body_new = ProtoField.bool("pictochat.body.new", "New")
 dissect_msg_type[MSG_TYPE.MSG_BODY] = function (buffer, pinfo, tree)
-	local payload
-	buffer, payload = dissect_payload(buffer, pinfo, tree)
+	local subtree = tree:add(p, buffer(), string.format("PictoChat Message Body: %d bytes", buffer:len()))
 
-	payload:add_le(p.fields.unknown, buffer(0x00, 2))
-	payload:add_le(p.fields.body_len, buffer(0x02, 1))
+	subtree:add_le(p.fields.unknown, buffer(0x00, 2))
+	subtree:add_le(p.fields.body_len, buffer(0x02, 1))
 	local body_len = buffer(0x02, 1):le_uint()
 
-	payload:add_le(p.fields.body_end, buffer(0x03, 1))
+	subtree:add_le(p.fields.body_end, buffer(0x03, 1))
 	-- This appears to be some kind of offset for indicating where to store the
 	-- current frame's data in a larger buffer. Messages sent in multiple parts
 	-- increment this value by 160 for each new original (p.fields.original ==
 	-- True) message.
-	payload:add_le(p.fields.body_offset, buffer(0x04, 2))
+	subtree:add_le(p.fields.body_offset, buffer(0x04, 2))
 	local body_offset = buffer(0x04, 2):le_uint()
-	payload:add_le(p.fields.unknown, buffer(0x06, 2)) -- usually 0
-
+	subtree:add_le(p.fields.unknown, buffer(0x06, 2)) -- usually 0
 	buffer = buffer(0x08)
+
 	-- This appears to be the actual message data (the drawing) sent as an
 	-- array of 8x8 tiles.
-	payload:add(p.fields.body, buffer(0, body_len))
-	if state.want_next == body_offset then
-		state.want_next = body_offset + body_len
-		state.body_remaining = state.body_remaining - body_len
-		state.buf:append(buffer(0, body_len:bytes()))
-		ByteArray.tvb(state.buf, "Complete message")
-	end
+	local body = buffer(0, body_len)
+	subtree:add(p.fields.body, body)
+
 	buffer = buffer(body_len)
 
-	payload:add_le(p.fields.body_seq, buffer(0x00, 2))
-	payload:add_le(p.fields.unknown, buffer(0x02, 2)) -- copy
+	subtree:add_le(p.fields.body_seq, buffer(0x00, 2))
+	subtree:add_le(p.fields.body_new, buffer(0x02, 2))
+	local new = buffer(0x02, 2):le_uint() ~= 0
+
+	-- reassembly
+	if pinfo.visited == false and state.msg[state.mid].buf:len() == body_offset and new == true then
+		state.msg[state.mid].buf:append(body:bytes())
+	end
+
+	state.pid_mid_map[pinfo.number] = state.mid
+
+	if pinfo.visited then
+		ByteArray.tvb(state.msg[state.pid_mid_map[pinfo.number]].buf, "Complete message")
+	end
 end
 
+p.fields.status_id = ProtoField.uint32("pictochat.status.id", "ID")
 p.fields.status_addr = ProtoField.bytes("pictochat.status.addr", "Address", base.COLON)
 p.fields.status_seq = ProtoField.uint16("pictochat.status.seq", "Sequence")
+p.fields.status_new = ProtoField.bool("pictochat.status.new", "New")
 dissect_msg_type[MSG_TYPE.STATUS] = function (buffer, pinfo, tree)
-	local payload
-	buffer, payload = dissect_payload(buffer, pinfo, tree)
+	local subtree = tree:add(p, buffer(), string.format("PictoChat Status: %d bytes", buffer:len()))
 
-	payload:add_le(p.fields.unknown, buffer(0x00, 4))
+	subtree:add_le(p.fields.status_id, buffer(0x00, 4))
+	local id = buffer(0x00, 4):le_uint()
 
 	buffer = buffer(0x04)
 
 	for _ = 0, 15 do
-		add_addr_le(payload, p.fields.status_addr, buffer(0, 6))
+		add_addr_le(subtree, p.fields.status_addr, buffer(0, 6))
 		buffer = buffer(6)
 	end
 
-	payload:add_le(p.fields.status_seq, buffer(0x00, 2))
-	payload:add_le(p.fields.unknown, buffer(0x02, 2))
+	subtree:add_le(p.fields.status_seq, buffer(0x00, 2))
+	subtree:add_le(p.fields.status_new, buffer(0x02, 2))
+
+	state.pid_mid_map[pinfo.number] = id
 end
+dissect_msg_type[MSG_TYPE.STATUS_END] = dissect_msg_type[MSG_TYPE.STATUS]
 
 p.fields.magic = ProtoField.bytes("pictochat.magic", "Magic")
 p.fields.new = ProtoField.bool("pictochat.new", "New")
+p.fields.flags = ProtoField.uint16("pictochat.flags", "Flags", base.HEX, nil, 0xffff)
 p.fields.msg_type = ProtoField.uint8("pictochat.msg_type", "Message type", base.DEC, MSG_TYPE_MAP)
+p.fields.len = ProtoField.uint16("pictochat.len", "Payload length")
+p.fields.data = ProtoField.bytes("pictochat.data", "Payload")
 function p.dissector(buffer, pinfo, tree)
 	-- check magic
 	if buffer(0x00, 2):uint() ~= GAMEID.PICTOCHAT then return 0 end
 	-- check length (avoid errors)
 	if buffer:len() < 6 then return 0 end
 
-	local subtree = tree:add(p, buffer(), string.format("%s: %d bytes", p.description, buffer():len()))
-	local buffer_len = buffer():len()
+	pinfo.cols.protocol = p.name
+	local buffer_len = math.min(buffer:len(), 0x0a)
+	local subtree = tree:add(p, buffer(0, buffer_len), string.format("PictoChat Header: %d bytes", buffer_len))
 
   subtree:add_le(p.fields.magic, buffer(0x00, 2))
   subtree:add_le(p.fields.new, buffer(0x02, 2))
-	subtree:add_le(p.fields.msg_type, buffer(0x04, 1))
-	local msg_type = buffer(0x04, 1):le_uint()
+	subtree:add_le(p.fields.flags, buffer(0x04, 2))
+	local flags = buffer(0x04, 2):le_uint()
+	if flags == 0 then
+		pinfo.cols.info = "???"
+		return buffer_len
+	end
+
+  subtree:add_le(p.fields.msg_type, buffer(0x06, 2))
+	local msg_type = buffer(0x06, 2):le_uint()
 	local msg_type_str = MSG_TYPE_MAP[msg_type]
-  subtree:add_le(p.fields.unknown, buffer(0x05, 1))
 
-	-- pretty wireshark shit
-	pinfo.cols.protocol = p.name
 	if msg_type_str ~= nil then
 		pinfo.cols.info = msg_type_str
 	end
 
-	if msg_type == MSG_TYPE.INIT then
-		-- no more content
-		return buffer_len
-	end
+	subtree:add_le(p.fields.len, buffer(0x08, 2))
+	local payload_length = buffer(0x08, 2):le_uint()
+  local payload = buffer(0x0a, payload_length)
 
-  subtree:add_le(p.fields.unknown, buffer(0x06, 2))
-	buffer = buffer(0x08)
+	subtree:add(p.fields.data, payload)
 
 	local subdissector = dissect_msg_type[msg_type]
 	if subdissector ~= nil then
-		subdissector(buffer, pinfo, subtree)
+		subdissector(payload, pinfo, tree)
 	end
 
+	-- prefix info field with message ID
+	pinfo.cols.info = string.format("[%08x] %s", state.pid_mid_map[pinfo.number] or 0, pinfo.cols.info)
+
 	return buffer_len
 end