-- [[ uncomment it to disable this module ]]--
-- do return end

g_mbtcp = Proto("ModbusTCP", "Modbus TCP")

print("Script ModbusTCP", NotWireshark)

function tableContains(tbl, key)
    if tbl == nil then
        return false
    end
    for k, v in pairs(tbl) do
        if k == key then
            return true
        end
    end

    return false
end

-- g_mbtcp.enable = false

local functions_code = {
    [1] = "Read Coils",
    [2] = "Read Discrete Inputs",
    [3] = "Read Holding Registers",
    [4] = "Read Input Registers",
    [5] = "Write Single Coil",
    [6] = "Write Holding Register",
    [0x0F] = "Write (N) Coils",
    [0x10] = "Write (N) Holding Registers",

    [0x81] = "Read Coils error",
    [0x82] = "Read Discrete Inputs error",
    [0x83] = "Read Holding Registers error",
    [0x84] = "Read Input Registers error",
    [0x85] = "Write Single Coil error",
    [0x86] = "Write Holding Register error",
    [0x8F] = "Write (N) Coils error",
    [0x90] = "Write (N) Holding Registers error",
}

local exception_code = {
    [1] = "Illegal function",
    [2] = "Illegal address",
    [3] = "Illegal data",
    [4] = "Slave device failure",
    [5] = "Ack",
    [6] = "Slave device busy",
    [8] = "Checksum error",
    [0xA] = "Gateway unavailable",
    [0XB] = "Gateway response failure"
}

local coils_code = {
    [0] = "Off",
    [0x00FF] = "On",
    [0xFF00] = "On"
}

-- define protocol fields
-- local f1 = ProtoField.new("mbtcp.BOOLEAN", "BOOLEAN", ftypes.BOOLEAN)
-- local pf_bool = ProtoField.bool("mbtcp.BOOLEAN", "BOOLEAN", base.NONE)--, base.HEX, {'True', 'False'}, 0xF0, 'blob?')

local pf_transId = ProtoField.uint16("mbtcp.TransactionId", "Transaction Identifier")
local pf_protoId = ProtoField.uint16("mbtcp.ProtocolId", "Protocol Identifier")
local pf_len = ProtoField.uint16("mbtcp.Length", "Length")
local pf_unitId = ProtoField.uint8("mbtcp.UnitId", "Unit Identifier")
local pf_fcode = ProtoField.uint8("mbtcp.Function", "Function", base.HEX, functions_code)

local pf_size = ProtoField.uint8("mbtcp.Size", "Size")
local pf_err = ProtoField.uint8("mbtcp.ErrCode", "ErrCode", base.HEX, exception_code)
local pf_addr = ProtoField.uint16("mbtcp.RegAddr", "Address")
local pf_cnt = ProtoField.uint16("mbtcp.RegCount", "Count")
local pf_Wcnt = ProtoField.uint16("mbtcp.WordCount", "Word Count")
local pf_Bcnt = ProtoField.uint8("mbtcp.ByteCount", "Byte Count")
local pf_bcnt = ProtoField.uint16("mbtcp.BitCount", "Bit Count")
local pf_reg = ProtoField.uint16("mbtcp.RegValue", "RegValue", base.HEX)
local pf_crc = ProtoField.uint16("mbtcp.Crc16", "Crc16", base.HEX)

local pf_bits = ProtoField.uint8("mbtcp.Bits", "Bits", base.HEX)
local pf_coil = ProtoField.uint8("mbtcp.Coils", "Coils", base.HEX, coils_code)


g_mbtcp.fields = {pf_transId, pf_protoId, pf_len, pf_unitId,
                pf_id,  pf_fcode, pf_size, pf_err,
                pf_addr,  pf_cnt, pf_Wcnt, pf_Bcnt, pf_bcnt, pf_reg,  pf_crc,
                pf_bits,  pf_coil
                }

-- define dissector function (Tvb, Pinfo, TreeItem)
function g_mbtcp.dissector(buf, pkt, tree)
    local pktlen = buf:reported_length_remaining()

    if pktlen < 8 then
        return 0
    end

    local length = buf:range(4, 2):uint()

    if length + 6 ~= pktlen then
        return 0
    end

    local fcode = buf(7, 1):uint()

    if not tableContains(functions_code, fcode) then
        return 0
    end

    local subtree = tree:add(g_mbtcp, buf:range(0, pktlen))

    subtree:add(pf_transId, buf:range(0, 2))
    subtree:add(pf_protoId, buf:range(2, 2))
    subtree:add(pf_len, buf:range(4, 2))
    subtree:add(pf_unitId, buf:range(6, 1))

    if (fcode & 0x80) == 0x80 then
        if pktlen == 5 then
            local subtree = tree:add(g_mbtcp, buf:range(0, pktlen))

            subtree:add(pf_id, buf:range(0, 1))
            subtree:add(pf_fcode, buf:range(1, 1))
            subtree:add(pf_err, buf:range(2, 1))
            subtree:add(pf_crc, crc)
        else
            return 0
        end
    elseif fcode == 1 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_bcnt, buf:range(10, 2))
        end

        local cnt = buf:range(8, 1):uint()
        if (cnt + 9) == pktlen then
            local subtree = tree:add(g_mbtcp, buf:range(0, pktlen))
            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_Bcnt, buf:range(8, 1))

            local off = 9
            repeat
                subtree:add(pf_bits, buf:range(off, 1))
                off = off + 1
                cnt = cnt - 1
            until cnt == 0
        end
    elseif fcode == 2 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_bcnt, buf:range(10, 2))
        end

        local cnt = buf:range(8, 1):uint()
        if (cnt + 9) == pktlen then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))
            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_Bcnt, buf:range(8, 1))

            local off = 9
            repeat
                subtree:add(pf_bits, buf:range(off, 1))
                off = off + 1
                cnt = cnt - 1
            until cnt == 0
        end
    elseif fcode == 3 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_Wcnt, buf:range(10, 2))
        else
            local cnt = buf:range(8, 1):uint()
            if (cnt % 2) == 1 then
                return 0
            end
            if (cnt + 9) == pktlen then
                local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

                subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
                subtree:add(pf_Bcnt, buf:range(8, 1))

                local off = 9
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0
            else
                return 0
            end
        end
    elseif fcode == 4 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_Wcnt, buf:range(10, 2))
        else
            local cnt = buf:range(8, 1):uint()
            if (cnt % 2) == 1 then
                return 0
            end
            if (cnt + 9) == pktlen then
                local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

                subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
                subtree:add(pf_Bcnt, buf:range(8, 1))

                local off = 9
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0
            else
                return 0
            end
        end
    elseif fcode == 5 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_coil, buf:range(10, 2))

            subtree = tree:add(g_mbtcp, buf:range(0, pktlen))
            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_coil, buf:range(10, 2))
        else
            return 0
        end
    elseif fcode == 6 then
        if pktlen == 8 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1))
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_reg, buf:range(10, 2))

            subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))
            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(8, 2))
            subtree:add(pf_reg, buf:range(10, 2))
        else
            return 0
        end
    elseif fcode == 0x0F then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(8, 2)):append_text(" <Start>")
            subtree:add(pf_bcnt, buf:range(10, 2))
        elseif pktlen >= 14 then
            local cnt = buf:range(10, 2):uint()
            if (cnt + 12) == pktlen then
                local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

                subtree:add(pf_fcode, buf:range(7, 1))
                subtree:add(pf_addr, buf:range(8, 2)):append_text(" <Start>")
                subtree:add(pf_bcnt, buf:range(10, 2))
                subtree:add(pf_Bcnt, buf:range(12, 1))

                local off = 13
                repeat
                    subtree:add(pf_bits, buf:range(off, 1))
                    off = off + 1
                    cnt = cnt - 1
                until cnt == 0
            else
                return 0
            end
        else
            return 0
        end
    elseif fcode == 0x10 then
        if pktlen == 12 then
            local subtree = tree:add(g_mbtcp, buf:range(7, 5))

            subtree:add(pf_fcode, buf:range(7, 1)):append_text(" <Respond>")
            subtree:add(pf_addr, buf:range(8, 2)):append_text(" <Start>")
            subtree:add(pf_Wcnt, buf:range(10, 2))
        elseif pktlen >= 17 then
            local Wcnt = buf:range(10, 2):uint()
            local Bcnt = buf:range(12, 1):uint()

            if (Wcnt * 2) ~= Bcnt then
                return 0
            end
            if (Bcnt + 13) ~= pktlen then
                return 0
            end

                local subtree = tree:add(g_mbtcp, buf:range(7, pktlen-7))

                subtree:add(pf_fcode, buf:range(7, 1))
                subtree:add(pf_addr, buf:range(8, 2)):append_text(" <Start>")
                subtree:add(pf_Wcnt, buf:range(10, 2))
                subtree:add(pf_Bcnt, buf:range(12, 1))

                local off = 13
                repeat
                    subtree:add(pf_reg, buf:range(off, 2))
                    off = off + 2
                    cnt = cnt - 2
                until cnt == 0
        else
            return 0    -- unknown
        end
    end


    print("Script ModbusTCP dissector exit")
end

if NotWireshark == nil then
    local udp_encap_table = DissectorTable.get("tcp.port")

    udp_encap_table:add(502, p_smsProto)
end

print("Script ModbusTCP exit")