-- [[ 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")