msgpack=require("e3pf-mp")

epf_protocol = Proto("e3pf", "e3team Protocol Framework")

packet_length = ProtoField.uint64("e3pf.packet_length", "Packet Length", base.DEC)
packet_id = ProtoField.uint8("e3pf.packet_id", "Packet ID", base.DEC)

-- Message: ClientHello
ch_protocol_version = ProtoField.uint32("e3pf.client_hello.protocol_version", "Protocol Version", base.DEC)
ch_client_random = ProtoField.string("e3pf.client_hello.client_random", "Client Random")
-- Client certificate: is a subtree
ch_client_x25519 = ProtoField.string("e3pf.client_hello.client_x25519", "Client X25519 Public Key")

-- Message: ServerHello
sh_protocol_version = ProtoField.uint32("e3pf.server_hello.protocol_version", "Protocol Version", base.DEC)
-- Server Certificate: is a subtree
sh_server_random = ProtoField.string("e3pf.server_hello.server_random", "Server Random")
sh_server_x25519 = ProtoField.string("e3pf.server_hello.server_x25519", "Server X25519 Public Key")

-- Struct: Certificate
-- Details: is a subtree
c_fingerprint = ProtoField.string("e3pf.certificate.fingerprint", "Fingerprint")
c_signature = ProtoField.string("e3pf.certificate.signature", "Signature")

-- Struct: Certificate.Details
c_d_name = ProtoField.string("e3pf.certificate.details.name", "Name")
c_d_not_before = ProtoField.string("e3pf.certificate.details.not_before", "Not Before")
c_d_not_after = ProtoField.string("e3pf.certificate.details.not_after", "Not After")
c_d_public_key = ProtoField.string("e3pf.certificate.details.public_key", "Public Key")
c_d_issuer_public_key = ProtoField.string("e3pf.certificate.details.issuer_public_key", "Issuer Public Key")
-- TODO: Claims

epf_protocol.fields = { packet_length, packet_id, ch_protocol_version, ch_client_random, ch_client_x25519, c_fingerprint, c_signature, c_d_name, c_d_not_before, c_d_not_after, c_d_public_key, c_d_issuer_public_key, sh_protocol_version, sh_server_random, sh_server_x25519 }

function dump(o)
   if type(o) == 'table' then
      local s = '{ '
      for k,v in pairs(o) do
         if type(k) ~= 'number' then k = '"'..k..'"' end
         s = s .. '['..k..'] = ' .. dump(v) .. ','
      end
      return s .. '} '
   else
      return tostring(o)
   end
end

function id_to_name(id)
	if id == 1 then return "ClientHello"
	elseif id == 2 then return "ServerHello"
	elseif id == 3 then return "HandshakeFinished"
	elseif id == 4 then return "ApplicationData"
	end
end

local function table_to_hex(ctable)
	local sb=''
	for _, value in pairs(ctable) do
		sb=sb .. string.format("%x", value)
	end
	return sb
end

local function cert_to_tree(rtree, name, ctable, buffer)
	local rrtree = rtree:add(epf_protocol, buffer(), name)
	
	local tree = rrtree:add(epf_protocol, buffer(), "Details")

	local details = ctable[1]

	tree:add(c_d_name, details[1])
	tree:add(c_d_not_before, tostring(details[2]))
	tree:add(c_d_not_after, tostring(details[3]))
	tree:add(c_d_public_key, table_to_hex(details[4]))
	tree:add(c_d_issuer_public_key, table_to_hex(details[5]))

	rrtree:add(c_fingerprint, ctable[2])
	rrtree:add(c_signature, table_to_hex(ctable[3]))
end

function epf_protocol.dissector(buffer, pinfo, tree)
	length = buffer:len()
	if length == 0 then return end

	pinfo.cols.protocol = epf_protocol.name

	local subtree = tree:add(epf_protocol, buffer(), "e3team Protocol Framework Data")
	local dtree = subtree:add(epf_protocol, buffer(), "Packet Data")

	local packet_length_flag = tonumber(buffer(0,8):le_uint64())

	subtree:add_le(packet_length, buffer(0, 8))

	local data = buffer(8, packet_length_flag)

	local msgpack_data=''
	local bytes_message = data:bytes()
	for i=0, (bytes_message:len()-1) do
		value=bytes_message:get_index(i)
		msgpack_data=msgpack_data .. string.char(value)
	end

	local message_unpacked = msgpack.unpack(msgpack_data)

	subtree:add(packet_id, message_unpacked[1]):append_text(" (" .. id_to_name(message_unpacked[1]) .. ")")

	local packet_bytes = message_unpacked[2]
	local packet_data=''
	for _, value in pairs(packet_bytes) do
		packet_data=packet_data .. string.char(value)
	end

	local packet = msgpack.unpack(packet_data)

	if message_unpacked[1] == 1 then
		-- CLIENT HELLO
		dtree:add(ch_protocol_version, packet[1])
		dtree:add(ch_client_random, table_to_hex(packet[2]))
		dtree:add(ch_client_x25519, table_to_hex(packet[4]))
	elseif message_unpacked[1] == 2 then
		-- SERVER HELLO
		dtree:add(sh_protocol_version, packet[1])
		cert_to_tree(dtree, "Server Certificate", packet[2], buffer)
		dtree:add(sh_server_random, table_to_hex(packet[3]))
		dtree:add(sh_server_x25519, table_to_hex(packet[4]))
	end
end

local function heuristic_checker(buffer, pinfo, tree)
	length = buffer:len()
	if length < 12 then return false end

	local header_start_flag = buffer(8,1):uint()

	if header_start_flag ~= 0x92 then return false end

	local packet_length_flag = buffer(0,8):le_uint64()

	local packet_length = length - 8

	if tostring(packet_length) ~= tostring(packet_length_flag) then return false end

	local packet_id = buffer(9,1):uint()

	print("packet " .. packet_id)

	if packet_id == 0x01 or packet_id == 0x02 or packet_id == 0x03 or packet_id == 0x04 then
		epf_protocol.dissector(buffer, pinfo, tree)
		return true
	else
		return false
	end
end

epf_protocol:register_heuristic("tcp", heuristic_checker)