-- Declare the custom protocol
myprotocol = Proto("myprotocol", "My Custom Protocol")

-- Define fields based on the header structure
local f_sequence_number = ProtoField.uint32("myprotocol.sequence_number", "Sequence Number", base.DEC)
local f_acknowledgment_number = ProtoField.uint32("myprotocol.acknowledgment_number", "Acknowledgment Number", base.DEC)
local f_flags = ProtoField.uint8("myprotocol.flags", "Flags", base.HEX)
local f_checksum = ProtoField.uint16("myprotocol.checksum", "Checksum", base.HEX)
local f_options = ProtoField.uint8("myprotocol.options", "Options", base.HEX)
local f_window = ProtoField.uint16("myprotocol.window", "Window Size", base.DEC)

-- Add fields to the protocol
myprotocol.fields = {f_sequence_number, f_acknowledgment_number, f_flags, f_checksum, f_options, f_window}

-- Define Flags
local SYN = 0x20    -- Synchronization
local ACK = 0x10    -- Acknowledgment
local PSH = 0x08    -- Push
local END = 0x04    -- End
local KP  = 0x02    -- Keep-Alive
local ERR = 0x01    -- Error
local SYN_ACK = SYN + ACK
local END_ACK = END + ACK
local KP_ACK = KP + ACK
local PSH_ACK = PSH + ACK
local PSH_ACK_ERR = PSH + ACK + ERR

-- Define Options
local INIT = 0x08   -- Initialization
local MF   = 0x04   -- More Fragments
local DF   = 0x02   -- Don't Fragment
local FIN  = 0x01   -- Finish

-- Function to interpret flags
local function dissect_flags(flags_value)
    local flag_map = {
        [SYN] = "SYN",
        [ACK] = "ACK",
        [SYN_ACK] = "SYN-ACK",
        [PSH] = "PSH",
        [END] = "END",
        [PSH_ACK] = "PSH-ACK",
        [END_ACK] = "END-ACK",
        [KP] = "KP",
        [KP_ACK] = "KP-ACK",
        [PSH_ACK_ERR] = "PSH-ACK-ERR",
        [ERR] = "ERR"
    }

    -- Default to "Unknown Flags" if flag is not found
    return flag_map[flags_value] or "Not set"
end

-- Function to interpret options
local function dissect_options(options_value)
    local options_map = {
        [INIT] = "INIT",
        [MF] = "MF",
        [DF] = "DF",
        [FIN] = "FIN"
    }

    -- Default to "Unknown Options" if option is not found
    return options_map[options_value] or "Not set"
end

-- Dissector function
function myprotocol.dissector(buffer, pinfo, tree)
    if buffer:len() < 14 then return end -- Ensure the packet is long enough for headers

    -- Protocol column
    pinfo.cols.protocol = myprotocol.name

    -- Create subtree
    local subtree = tree:add(myprotocol, buffer(), "My Protocol Data")

  -- Sequence Number (Decimal only)
    local sequence_number = buffer(0, 4):uint() -- Use uint() for big-endian or le_uint() for little-endian
    subtree:add(f_sequence_number, buffer(0, 4)):set_text(string.format("Sequence Number: %d", sequence_number))

    -- Acknowledgment Number (Decimal only)
    local acknowledgment_number = buffer(4, 4):uint() -- Use uint() for big-endian or le_uint() for little-endian
    subtree:add(f_acknowledgment_number, buffer(4, 4)):set_text(string.format("Acknowledgment Number: %d", acknowledgment_number))

    -- Flags
    local flags_value = buffer(8, 1):uint()
    local flags_string = dissect_flags(flags_value)
    subtree:add(f_flags, buffer(8, 1)):append_text(" [" .. flags_string .. "]")

    -- Checksum
    local checksum_value = buffer(9, 2):uint()
    subtree:add(f_checksum, buffer(9, 2)):set_text(string.format("Checksum: 0x%04X", checksum_value))

    -- Options
    local options_value = buffer(11, 1):uint()
    local options_string = dissect_options(options_value)
    subtree:add(f_options, buffer(11, 1)):append_text(" [" .. options_string .. "]")

    -- Window Size
    local window_size = buffer(12, 2):le_uint()
    subtree:add(f_window, buffer(12, 2)):set_text(string.format("Window Size: %d", window_size))

    -- Set Info column
    pinfo.cols.info = string.format("Sequence number: %d, Acknowledgment number: %d, Flags: %s, Options: %s", sequence_number, acknowledgment_number, flags_string, options_string)
end

-- Register dissector on specific UDP ports
local udp_table = DissectorTable.get("udp.port")
udp_table:add(8080, myprotocol)
udp_table:add(8000, myprotocol)
udp_table:add(7000, myprotocol)