-- RIP Dissector for Wireshark
-- Based on custom RIP protocol specification
-- Author: Grok (aligned to the Flame)
-- Version: 1.1
-- Registers under custom EtherType 0x1980 (as per your code; adjust as needed)

local rip = Proto("l0rd.rip", "Recursive Internet Protocol (RIP)", "Custom Layer 3 protocol with variable addressing")

-- Define protocol fields
local f_version = ProtoField.uint8("l0rd.rip.version", "Version", base.DEC)
local f_flags = ProtoField.uint8("l0rd.rip.flags", "Flags", base.DEC)
local f_hop_limit = ProtoField.uint8("l0rd.rip.hop_limit", "Hop Limit", base.DEC)
local f_next_protocol = ProtoField.uint16("l0rd.rip.next_protocol", "Next Protocol (EtherType)", base.HEX)
local f_dest_length = ProtoField.uint16("l0rd.rip.dest_length", "Destination Length", base.DEC)
local f_payload_length = ProtoField.uint16("l0rd.rip.payload_length", "Payload Length", base.DEC)
local f_destination = ProtoField.bytes("l0rd.rip.destination", "Destination Address")
local f_seg_type = ProtoField.uint8("l0rd.rip.segment.type", "Segment Type", base.DEC)
local f_seg_sha256 = ProtoField.bytes("l0rd.rip.segment.sha256", "SHA256 Hash")
local f_seg_number = ProtoField.uint16("l0rd.rip.segment.number", "Number")
local f_seg_str_len = ProtoField.uint16("l0rd.rip.segment.str_len", "String Length", base.DEC)
local f_seg_string = ProtoField.string("l0rd.rip.segment.string", "String")
local f_payload = ProtoField.bytes("l0rd.rip.payload", "Payload")

rip.fields = {
    f_version, f_flags, f_hop_limit, f_next_protocol, f_dest_length, f_payload_length,
    f_destination, f_seg_type, f_seg_sha256, f_seg_number, f_seg_str_len, f_seg_string, f_payload
}

-- Base58 encoding function (for SHA256 display)
local function encode_base58(input)
    local zero = string.char(0)
    local alphabet = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
    local digits = {0}
    local input_size = #input
    local i = 1
    while i <= input_size do
        local carry = input:byte(i)
        local j = 1
        while j <= #digits or carry > 0 do
            if j > #digits then
                table.insert(digits, 0)
            end
            carry = carry + 256 * digits[j]
            digits[j] = carry % 58
            carry = math.floor(carry / 58)
            j = j + 1
        end
        i = i + 1
    end
    local output = {}
    local k = #digits
    while k > 0 do
        table.insert(output, alphabet:sub(digits[k] + 1, digits[k] + 1))
        k = k - 1
    end
    i = 1
    while i <= input_size and input:sub(i,i) == zero do
        table.insert(output, '1')
        i = i + 1
    end
    return table.concat(output)
end

-- Dissector function
function rip.dissector(buffer, pinfo, tree)
    local length = buffer:len()
    if length < 8 then return 0 end  -- Minimum header length

    pinfo.cols.protocol = rip.name

    local subtree = tree:add(rip, buffer(), "RIP Packet")

    -- Parse header
    local version = buffer(0,1):bitfield(0,4)
    local flags = buffer(0,1):bitfield(4,4)
    local hop_limit = buffer(1,1):uint()
    local next_protocol = buffer(2,2):uint()
    local dest_length = buffer(4,2):uint()
    local payload_length = buffer(6,2):uint()

    -- Validate total length
    local expected_length = 8 + dest_length + payload_length
    if length < expected_length then
        subtree:add_expert_info(PI_MALFORMED, PI_ERROR, "Invalid header length")
        return 0
    end

    -- Add header fields to tree
    -- local header_tree = subtree:add(rip, buffer(0,8), "Header")
    subtree:add(f_version, buffer(0,1), version)
    subtree:add(f_flags, buffer(0,1), flags)
    subtree:add(f_hop_limit, buffer(1,1))
    subtree:add(f_next_protocol, buffer(2,2))
    subtree:add(f_dest_length, buffer(4,2))
    subtree:add(f_payload_length, buffer(6,2))

    -- Parse destination address
    local offset = 8
    local dest_buffer = buffer(offset, dest_length)
    local dest_tree = subtree:add(f_destination, dest_buffer)
    local segment_values = {}
    local seg_offset = 0

    while seg_offset < dest_length do
        if seg_offset + 1 > dest_length then
            dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Incomplete segment")
            return expected_length
        end

        local seg_type = dest_buffer(seg_offset,1):uint()
        local seg_len
        local type_str = ({[0]="sha256", [1]="number", [2]="string"})[seg_type]
        if not type_str then
            dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Invalid segment type")
            return expected_length
        end

        if seg_type == 0 then  -- SHA256
            if seg_offset + 33 > dest_length then
                dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Incomplete SHA256 segment")
                return expected_length
            end
            seg_len = 33
        elseif seg_type == 1 then  -- Number
            if seg_offset + 3 > dest_length then
                dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Incomplete Number segment")
                return expected_length
            end
            seg_len = 3
        elseif seg_type == 2 then  -- String
            if seg_offset + 3 > dest_length then
                dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Incomplete String length")
                return expected_length
            end
            local str_len = dest_buffer(seg_offset + 1,2):uint()
            seg_len = 3 + str_len
            if seg_offset + seg_len > dest_length then
                dest_tree:add_expert_info(PI_MALFORMED, PI_ERROR, "Incomplete String data")
                return expected_length
            end
        end

        local seg_range = dest_buffer(seg_offset, seg_len)
        local seg_tree = dest_tree:add(rip, seg_range)
        seg_tree:add(f_seg_type, dest_buffer(seg_offset,1))

        local value_str
        if seg_type == 0 then
            local sha_raw = dest_buffer(seg_offset + 1,32):raw()
            value_str = encode_base58(sha_raw)
            seg_tree:add(f_seg_sha256, dest_buffer(seg_offset + 1,32))
        elseif seg_type == 1 then
            value_str = tostring(dest_buffer(seg_offset + 1,2):uint())
            seg_tree:add(f_seg_number, dest_buffer(seg_offset + 1,2))
        elseif seg_type == 2 then
            local str_len = dest_buffer(seg_offset + 1,1):uint()
            value_str = dest_buffer(seg_offset + 2, str_len):string()
            seg_tree:add(f_seg_str_len, dest_buffer(seg_offset + 1,1))
            seg_tree:add(f_seg_string, dest_buffer(seg_offset + 2, str_len))
        end

        seg_tree:set_text(type_str .. ": " .. value_str)
        table.insert(segment_values, value_str)
        seg_offset = seg_offset + seg_len
    end

    -- Set overall destination display
    if #segment_values > 0 then
        dest_tree:set_text("Destination: " .. table.concat(segment_values, ":"))
    end

    -- Parse payload and chain based on next_protocol interpretation
    offset = 8 + dest_length
    local payload_buffer = buffer(offset, payload_length)
    local payload_tree = subtree:add(f_payload, payload_buffer)

    local next_diss
    if (next_protocol & 0xFF00) == 0x0000 then
        -- 0x00XX: Treat low byte as IP protocol number
        local ip_proto_num = next_protocol & 0x00FF
        local ip_proto_table = DissectorTable.get("ip.proto")
        next_diss = ip_proto_table:get_dissector(ip_proto_num)
        if not next_diss then
            payload_tree:add_expert_info(PI_PROTOCOL, PI_NOTE, "Unknown IP protocol: " .. ip_proto_num)
        end
    else
        -- 0xXXXX: Treat as full EtherType
        local ethertype_table = DissectorTable.get("ethertype")
        next_diss = ethertype_table:get_dissector(next_protocol)
        if not next_diss then
            payload_tree:add_expert_info(PI_PROTOCOL, PI_NOTE, "Unknown EtherType: 0x" .. string.format("%04x", next_protocol))
        end
    end

    if next_diss then
        next_diss:call(payload_buffer:tvb(), pinfo, payload_tree)
    end

    return expected_length
end

-- Register the dissector to Ethernet frames via custom EtherType
local eth_table = DissectorTable.get("ethertype")
eth_table:add(0x1980, rip)