-- Example Custom Protocol Dissector for Wireshark/PyShark
-- Place this file in your Wireshark plugins directory:
--   Linux/Mac: ~/.local/lib/wireshark/plugins/ or ~/.wireshark/plugins/
--   Windows: %APPDATA%\Wireshark\plugins\
--
-- This dissector will automatically be available to PyShark via tshark

-- Create protocol
local custom_proto = Proto("custom", "Custom Airstream Protocol")

-- Define fields
local fields = {
    magic = ProtoField.uint32("custom.magic", "Magic Number", base.HEX),
    version = ProtoField.uint8("custom.version", "Version", base.DEC),
    msg_type = ProtoField.uint8("custom.msg_type", "Message Type", base.DEC, {
        [1] = "Data",
        [2] = "Control", 
        [3] = "Status",
        [4] = "Telemetry"
    }),
    length = ProtoField.uint16("custom.length", "Payload Length", base.DEC),
    sequence = ProtoField.uint32("custom.sequence", "Sequence Number", base.DEC),
    timestamp = ProtoField.uint64("custom.timestamp", "Timestamp", base.DEC),
    payload = ProtoField.bytes("custom.payload", "Payload")
}

-- Add fields to protocol
custom_proto.fields = fields

-- Dissector function
function custom_proto.dissector(buffer, pinfo, tree)
    -- Validate minimum length
    if buffer:len() < 20 then
        return
    end
    
    -- Check magic number (example: 0xABCD1234)
    local magic = buffer(0,4):uint()
    if magic ~= 0xABCD1234 then
        return  -- Not our protocol
    end
    
    -- Set protocol column in packet list
    pinfo.cols.protocol:set("CUSTOM")
    
    -- Create subtree for our protocol
    local subtree = tree:add(custom_proto, buffer(), "Custom Airstream Protocol")
    
    -- Add fields to tree
    subtree:add(fields.magic, buffer(0,4))
    subtree:add(fields.version, buffer(4,1))
    
    local msg_type = buffer(5,1):uint()
    subtree:add(fields.msg_type, buffer(5,1))
    
    subtree:add(fields.length, buffer(6,2))
    subtree:add(fields.sequence, buffer(8,4))
    subtree:add(fields.timestamp, buffer(12,8))
    
    -- Add payload if present
    local payload_len = buffer(6,2):uint()
    if buffer:len() >= 20 + payload_len then
        subtree:add(fields.payload, buffer(20, payload_len))
    end
    
    -- Update info column with summary
    local type_names = {[1]="Data", [2]="Control", [3]="Status", [4]="Telemetry"}
    local type_name = type_names[msg_type] or "Unknown"
    pinfo.cols.info:set(string.format("%s (Seq=%d, Len=%d)", 
                                      type_name, 
                                      buffer(8,4):uint(),
                                      payload_len))
end

-- Register dissector for specific UDP port
local udp_port = DissectorTable.get("udp.port")
udp_port:add(9999, custom_proto)  -- Listen on UDP port 9999

-- Also register as heuristic dissector for UDP
local function heuristic_checker(buffer, pinfo, tree)
    -- Check if this might be our protocol
    if buffer:len() < 20 then
        return false
    end
    
    -- Check magic number
    local magic = buffer(0,4):uint()
    if magic == 0xABCD1234 then
        custom_proto.dissector(buffer, pinfo, tree)
        return true
    end
    
    return false
end

-- Register heuristic dissector
custom_proto:register_heuristic("udp", heuristic_checker)