------------------------------------------------------------------------------
-- TCP Error Tracker                                                        --
--                                                                          --
-- Counts the number of TCP errors by TCP stream, and source and            --
--  destination IP addresses. Also supports Check Point Orchestrators when  --
--  used with the cp_mho.lua plugin - but not required.                     --
--                                                                          --
------------------------------------------------------------------------------
--                                                                          --
-- Copyright: Allan Que, 2023                                               --
-- License  : GPL-3.0-or-later                                              --
--            see http://www.gnu.org/licenses/gpl.html                      --
--                                                                          --
------------------------------------------------------------------------------
--                                                                          --
-- Copy file to your Personal Plugins directory                             --
--  See Help > About > Folders for paths                                    --
--                                                                          --
------------------------------------------------------------------------------

local function tcp_errors_tap()
    local title_s="TCP Analysis Errors"
    local mho_mode_b=false

    local stream_m = {}
    local src_m = {}
    local dst_m = {}
    local mho_m = {}
    
    -- Exclude packets from counting towards TCP errors
    local tap_filter_s=
        "tcp.analysis.flags"..
        "&& !tcp.analysis.keep_alive"..
        "&& !tcp.analysis.keep_alive_ack"

    -- Include any filters user has already set in Wireshark
    if get_filter() ~= "" then tap_filter_s=tap_filter_s.."&&"..get_filter(s) end

    -- Additional filters when operating with Check Point MHO g_tcpdump packets
    -- Accomodates for traffic thru one MHO and response thru another,
    --   and gets around capture timing issues that show up as out-of-order
    if get_preference("mho.enabled") then
        tap_filter_s=tap_filter_s..
            "&& !tcp.analysis.out_of_order"..
            "&& !tcp.analysis.ack_lost_segment"..
            "&& !tcp.analysis.spurious_retransmission"
        title_s=title_s.." - MHO Mode"
        mho_mode_b=true
    end

    -- Create the tcp tap and window
    local tap = Listener.new("tcp", tap_filter_s)
    local win = TextWindow.new(title_s)

    -- remove() must be after tap object is instantiated
    win:set_atclose(function ()
        -- Remove tap when not in use
        -- Must be outside tap.packet()
        tap:remove()
    end)

    local function increment_counter(field_o, table_m)
        local field_s = tostring(field_o)
        local count_i = table_m[field_s] or 0
        table_m[field_s] = count_i + 1
    end

    local function mho_id(abs_t)
        -- Extract subseconds and convert to integer with rounding
        local subsec_f = tonumber(abs_t) % 1 * 1000000
        local subsec_i = math.floor(subsec_f + 0.5)

        -- Chassis and Blade calculations are from
        --   $SMODIR/bin/gtcpdump-helper.sh
        local mho_ch = bit.rshift(bit.band(subsec_i, 16), 4) + 1
        local mho_bl = bit.band(subsec_i, 15)

        return mho_ch.."_"..mho_bl
    end

    -- function called for every single packet
    function tap.packet(pinfo, buffer, tap_data)
        increment_counter(tap_data.th_stream, stream_m)
        increment_counter(pinfo.src, src_m)
        increment_counter(pinfo.dst, dst_m)
        if mho_mode_b then
            increment_counter(mho_id(pinfo.abs_ts), mho_m)
        end
    end

    -- Output supplied table data to text window
    local function output_table(title_s, table_m, col1_width_i)
        win:append("** "..title_s.." **\n")
        for key_s, value_i in pairs(table_m) do
            win:append(string.format("%"..col1_width_i.."s - %5d\n",key_s, value_i))
        end
        win:append("\n")
    end

    -- Update text window with latest info - run periodically
    function tap.draw()
        win:clear()
        win:append(title_s.."\n")
        output_table("By Stream", stream_m, 5)
        output_table("By Source IP", src_m, 15)
        output_table("By Destination IP", dst_m, 15)
        if mho_mode_b then
            output_table("By MHO Id", mho_m, 4)
        end
    end

    -- Called at the end of capture run
    function tap.reset()
        win:clear()
        stream_m = {}
        src_m = {}
        dst_m = {}
    end

    -- Sets display filter to supplied string
    local function display_filter(filter_s)
        set_filter(filter_s)
        apply_filter()
    end

    -- Dialog boxes for Stream, IP and MHO buttons
    local function stream_dlg()
        new_dialog("Enter a Stream Number",
            function(stream_s)
                -- Validate input is a number
                if tonumber(stream_s) == nil then return end
                display_filter("tcp.stream=="..stream_s)
            end,
        "Stream")
    end

    local function ip_dlg()
        new_dialog("Enter an IPv4 Address",
            function(input_ip_s)
                _, _, valid_ip_s = string.find(input_ip_s, "^%s*(%d+%.%d+%.%d+%.%d+)%s*$")
                if valid_ip_s == nil then return end
                display_filter("ip.addr=="..valid_ip_s)
            end,
        "IPv4")
    end

    local function mho_dlg()
        new_dialog("Enter MHO Id",
            function(input_mho_s)
                _, _, valid_mho_s = string.find(input_mho_s, "^%s*([12]_1?%d)%s*$")
                if valid_mho_s == nil then return end
                display_filter("mho.id=="..valid_mho_s)
            end,
        "MHO Id")
    end

    -- Buttons at the bottom of text window
    -- Order = top-down is left-to-right
    win:add_button("Copy All", function()
        copy_to_clipboard(win:get_text())
    end)
    if mho_mode_b then win:add_button("MHO", mho_dlg) end
    win:add_button("IPv4", ip_dlg)
    win:add_button("Stream", stream_dlg)
    win:add_button("No Filter", function()
        display_filter("")
    end)

    -- Start sending packets into tap
    retap_packets()
end

-- Require GUI and versions 3.5 and later
if not (gui_enabled() and get_preference) then return end

register_menu("TCP/Error Tracker", tcp_errors_tap, MENU_TOOLS_UNSORTED)