//--------------------------------------------------------------------------
// Copyright (C) 2014-2026 Cisco and/or its affiliates. All rights reserved.
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License Version 2 as published
// by the Free Software Foundation.  You may not use, modify or distribute
// this program under any other version of the GNU General Public License.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
//--------------------------------------------------------------------------

// stream_module.cc author Russ Combs <rucombs@cisco.com>

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "stream_module.h"

#include "control/control.h"
#include "detection/rules.h"
#include "flow/dump_flows_descriptor.h"
#include "flow/dump_flows.h"
#include "flow/flow_cache.h"
#include "log/messages.h"
#include "lua/lua.h"
#include "main/analyzer_command.h"
#include "main/snort.h"
#include "main/snort_config.h"
#include "managers/inspector_manager.h"
#include "managers/module_manager.h"
#include "stream/flush_bucket.h"
#include "stream/tcp/tcp_stream_tracker.h"
#include "time/packet_time.h"
#include "trace/trace.h"

using namespace snort;
using namespace std;

#ifdef DEBUG_MSGS
static const TraceOption stream_trace_options[] =
{
    { "base", TRACE_BASE, "enable base stream trace logging" },
    { "flow", TRACE_FLOW, "enable flow trace logging" },
    { nullptr, 0, nullptr }
};
#endif

THREAD_LOCAL const Trace* stream_trace = nullptr;
static uint32_t hold_time = 0;

//-------------------------------------------------------------------------
// stream module
//-------------------------------------------------------------------------

#define FLOW_TYPE_PARAMS(name, idle) \
static const Parameter name[] = \
{ \
    { "idle_timeout", Parameter::PT_INT, "1:max32", idle, \
      "maximum inactive time before retiring session tracker" }, \
 \
    { nullptr, Parameter::PT_MAX, nullptr, nullptr, nullptr } \
}

static const Parameter allowlist_cache_params[] =
{
    { "enable", Parameter::PT_BOOL, nullptr, "false",
      "enable allowlist cache" },

    { "move_on_excess", Parameter::PT_BOOL, nullptr, "false",
      "move flows to allowlist instead of removing when max flows limit reached" },

    { nullptr, Parameter::PT_MAX, nullptr, nullptr, nullptr }
};

FLOW_TYPE_PARAMS(ip_params, "180");
FLOW_TYPE_PARAMS(icmp_params, "180");
FLOW_TYPE_PARAMS(tcp_params, "3600");
FLOW_TYPE_PARAMS(udp_params, "180");
FLOW_TYPE_PARAMS(user_params,"180");
FLOW_TYPE_PARAMS(file_params, "180");

#define FLOW_TYPE_TABLE(flow_type, proto, params) \
    { flow_type, Parameter::PT_TABLE, params, nullptr, \
      "configure " proto " cache limits" }

static const Parameter s_params[] =
{
#ifdef REG_TEST
    { "footprint", Parameter::PT_INT, "0:max32", "0",
      "use zero for production, non-zero for testing at given size (for TCP and user)" },
#endif

    { "held_packet_timeout", Parameter::PT_INT, "1:max32", "1000",
      "timeout in milliseconds for held packets" },

    { "ip_frags_only", Parameter::PT_BOOL, nullptr, "false",
      "don't process non-frag flows" },

    { "max_flows", Parameter::PT_INT, "2:max32", "476288",
      "maximum simultaneous flows tracked before pruning" },

    { "prune_flows", Parameter::PT_INT, "1:max32", "10",
      "maximum flows to prune at one time" },

    { "pruning_timeout", Parameter::PT_INT, "1:max32", "30",
      "minimum inactive time before being eligible for pruning" },

    { "require_3whs", Parameter::PT_INT, "-1:max31", "-1",
      "don't track midstream TCP sessions after given seconds from start up; -1 tracks all" },

    { "allowlist_cache", Parameter::PT_TABLE, allowlist_cache_params, nullptr, "configure allowlist cache" },

    { "drop_stale_packets", Parameter::PT_BOOL, nullptr, "false",
      "enable dropping of packets with stale timestamp" },

    FLOW_TYPE_TABLE("ip_cache",   "ip",   ip_params),
    FLOW_TYPE_TABLE("icmp_cache", "icmp", icmp_params),
    FLOW_TYPE_TABLE("tcp_cache",  "tcp",  tcp_params),
    FLOW_TYPE_TABLE("udp_cache",  "udp",  udp_params),
    FLOW_TYPE_TABLE("user_cache", "user", user_params),
    FLOW_TYPE_TABLE("file_cache", "file", file_params),

    { nullptr, Parameter::PT_MAX, nullptr, nullptr, nullptr }
};

// FIXIT-L setup and clear should extend to non-tcp flows as well
static const RuleMap stream_rules[] =
{
    { SESSION_EVENT_SYN_RX, "TCP SYN received" },
    { SESSION_EVENT_SETUP,  "TCP session established" },
    { SESSION_EVENT_CLEAR,  "TCP session cleared" },

    { 0, nullptr }
};

static const char* const flow_type_names[] =
{ "none", "ip_cache", "tcp_cache", "udp_cache", "icmp_cache", "user_cache", "file_cache", "max"};

StreamModule::StreamModule() : Module(MOD_NAME, MOD_HELP, s_params)
{ }

void StreamModule::set_trace(const Trace* trace) const
{ stream_trace = trace; }

const TraceOption* StreamModule::get_trace_options() const
{
#ifndef DEBUG_MSGS
    return nullptr;
#else
    return stream_trace_options;
#endif
}

enum DumpFlowsParams {FileName, Count, Protocol, ipA, ipB, portA, portB, Resume};
struct DumpFlowsLuaParameters
{
    DumpFlowsParams param;
    unsigned lua_index;
};

std::vector<DumpFlowsLuaParameters> df_params =
{
    { DumpFlowsParams::FileName, 1 },
    { DumpFlowsParams::Count, 2 },
    { DumpFlowsParams::Protocol, 3 },
    { DumpFlowsParams::ipA, 4 },
    { DumpFlowsParams::ipB, 5 },
    { DumpFlowsParams::portA, 6 },
    { DumpFlowsParams::portB, 7 },
    { DumpFlowsParams::Resume, 8 }
};

std::vector<DumpFlowsLuaParameters> dfb_params =
{
    { DumpFlowsParams::FileName, 1 },
    { DumpFlowsParams::Protocol, 2 },
    { DumpFlowsParams::ipA, 3 },
    { DumpFlowsParams::ipB, 4 },
    { DumpFlowsParams::portA, 5 },
    { DumpFlowsParams::portB, 6 },
    { DumpFlowsParams::Resume, 7 }
};

std::vector<DumpFlowsLuaParameters> dfs_params =
{
    { DumpFlowsParams::Protocol, 1 },
    { DumpFlowsParams::ipA, 2 },
    { DumpFlowsParams::ipB, 3 },
    { DumpFlowsParams::portA, 4 },
    { DumpFlowsParams::portB, 5 },
};

static int validate_cmd_parameters(ControlConn* ctrlcon, lua_State* L, std::vector<DumpFlowsLuaParameters>& cmd_params, DumpFlowsFilter* dff)
{
    for (auto const& param : cmd_params)
    {
        switch ( param.param )
        {
            case DumpFlowsParams::FileName:
            {
                const char* file_name = luaL_optstring(L, param.lua_index, nullptr);
                if ( !file_name )
                {
                    LogRespond(ctrlcon, "Dump flows requires a file name\n");
                    return false;
                }
                else
                    dff->file_name = file_name;

                break;
            }

            case DumpFlowsParams::Count:
                dff->count = luaL_optint(L, param.lua_index, 4);
                if ( dff->count >  1000 )
                {
                    LogRespond(ctrlcon, "Invalid count value: %u.  Count value must be between 1 - 1000\n", dff->count);
                    return false;
                }
                break;

            case DumpFlowsParams::Protocol:
            {
                const char* protocol = luaL_optstring(L, param.lua_index, nullptr);
                if ( !protocol )
                {
                    LogRespond(ctrlcon, "protocol must be a string or convertible to a string\n");
                    return false;
                }

                if ( *protocol )
                {
                    auto proto_it = protocol_to_type.find(protocol);
                    if ( proto_it == protocol_to_type.end() )
                    {
                        LogRespond(ctrlcon, "Invalid protocol: %s.  Valid protocols are IP/TCP/UDP/ICMP\n",
                            protocol);
                        return false;
                    }
                    else
                        dff->proto_type = proto_it->second;
                }

                break;
            }
            
            case DumpFlowsParams::ipA:
            {
                const char* ip_addr = luaL_optstring(L, param.lua_index, nullptr);
                if ( ip_addr  and !dff->set_ipA(ip_addr) )
                {
                    LogRespond(ctrlcon, "Invalid source ip\n");
                    return false;
                }

                break;
            }

            case DumpFlowsParams::ipB:
            {
                const char* ip_addr = luaL_optstring(L, param.lua_index, nullptr);
                if ( ip_addr and !dff->set_ipB(ip_addr) )
                {
                    LogRespond(ctrlcon, "Invalid destination ip\n");
                    return false;
                }

                break;
            }

            case DumpFlowsParams::portA:
            {
                unsigned port = luaL_optint(L, param.lua_index, 0);
                if ( port > 65535 )
                {
                    LogRespond(ctrlcon, "source_port must be between 0 - 65535\n");
                    return false;
                }
                if ( port )
                    dff->set_portA(static_cast<uint16_t>(port));

                break;
            }

            case DumpFlowsParams::portB:
            {
                unsigned port = luaL_optint(L, param.lua_index, 0);
                if ( port > 65535 )
                {
                    LogRespond(ctrlcon, "destination_port must be between 0 - 65535\n");
                    return false;
                }

                if ( port )
                    dff->set_portB(static_cast<uint16_t>(port));
                break;
            }
            
            case DumpFlowsParams::Resume:
                dff->resume = -1;
#ifdef REG_TEST
                dff->resume = luaL_optint(L, param.lua_index, -1);
#endif
                break;
        }
    }

    return true;   
}

static int dump_flows_worker(lua_State* L, bool enable_binary)
{
    ControlConn* ctrlcon = ControlConn::query_from_lua(L);
    Inspector* inspector = InspectorManager::get_inspector("stream", Module::GLOBAL);

    if (!inspector)
    {
        LogRespond(ctrlcon, "Dump flows requires stream to be configured\n");
        return -1;
    }

    DumpFlowsFilter* dff = nullptr;
    if ( enable_binary )
    {
        dff = new DumpFlowsFilterOr(enable_binary);
        if ( !validate_cmd_parameters(ctrlcon, L, dfb_params, dff) )
        {
            delete dff;
            return -1;
        }
    }
    else
    {
        dff = new DumpFlowsFilterAnd(enable_binary);
        if ( !validate_cmd_parameters(ctrlcon, L, df_params, dff) )
        {
            delete dff;
            return -1;
        }
    }

    // set filter for flow selection
    DumpFlows* df = new DumpFlows(ctrlcon, dff);

    LogRespond(ctrlcon, "== dumping connections\n");
    main_broadcast_command(df, ctrlcon);
    return 0;
}

static int dump_flows(lua_State* L)
{
    return dump_flows_worker(L, false);
}

static int dump_flows_binary(lua_State* L)
{
    return dump_flows_worker(L, true);
}
    
static int dump_flows_summary(lua_State* L)
{
    ControlConn* ctrlcon = ControlConn::query_from_lua(L);
    Inspector* inspector = InspectorManager::get_inspector("stream", Module::GLOBAL);

    if (!inspector)
    {
        LogRespond(ctrlcon, "Dump flows requires stream to be configured\n");
        return -1;
    }
    
    DumpFlowsFilter* dff = new DumpFlowsFilterAnd(false);
    if ( !validate_cmd_parameters(ctrlcon, L, dfs_params, dff) )
    {
        delete dff;
        return -1;
    }
    DumpFlowsSummary* dfs = new DumpFlowsSummary(ctrlcon, dff);

    LogRespond(ctrlcon, "== dumping connection summaries\n");
    main_broadcast_command(dfs, ctrlcon);
    return 0;
}

static const Command stream_cmds[] =
{
    { "dump_flows", dump_flows, nullptr, "dump the flow table in text format" },
    { "dump_flows_binary", dump_flows_binary, nullptr, "dump the flow table in binary format" },
    { "dump_flows_summary", dump_flows_summary, nullptr, "dump flow table summary" },
    { nullptr, nullptr, nullptr, nullptr }
};

const snort::Command* StreamModule::get_commands() const
{ return stream_cmds; }

const PegInfo* StreamModule::get_pegs() const
{ return base_pegs; }

PegCount* StreamModule::get_counts() const
{ return (PegCount*)&stream_base_stats; }

ProfileStats* StreamModule::get_profile() const
{ return &s5PerfStats; }

unsigned StreamModule::get_gid() const
{ return GID_SESSION; }

const RuleMap* StreamModule::get_rules() const
{ return stream_rules; }

const StreamModuleConfig* StreamModule::get_data()
{
    hold_time = config.hold_time;
    return &config;
}

uint32_t StreamModule::get_hold_time()
{ return hold_time; }

bool StreamModule::begin(const char* fqn, int, SnortConfig*)
{
    if ( !strcmp(fqn, MOD_NAME) )
        config = {};

    return true;
}

bool StreamModule::set(const char* fqn, Value& v, SnortConfig* c)
{
#ifdef REG_TEST
    if ( v.is("footprint") )
        config.footprint = v.get_uint32();
    else
#endif

    if ( v.is("held_packet_timeout") )
        config.hold_time = v.get_int32();

    else if ( v.is("ip_frags_only") )
    {
        if ( v.get_bool() )
            c->set_run_flags(RUN_FLAG__IP_FRAGS_ONLY);
    }
    else if ( v.is("max_flows") )
        config.flow_cache_cfg.max_flows = v.get_uint32();

    else if ( v.is("prune_flows") )
        config.flow_cache_cfg.prune_flows = v.get_uint32();

    else if ( v.is("pruning_timeout") )
        config.flow_cache_cfg.pruning_timeout = v.get_uint32();

    else if ( v.is("require_3whs") )
        config.hs_timeout = v.get_int32();

    else if ( !strcmp(fqn, "stream.allowlist_cache.enable") )
        config.flow_cache_cfg.allowlist_cache = v.get_bool();

    else if ( !strcmp(fqn, "stream.allowlist_cache.move_on_excess") )
        config.flow_cache_cfg.move_to_allowlist_on_excess = v.get_bool();

    else if ( !strcmp(fqn, "stream.file_cache.idle_timeout") )
        config.flow_cache_cfg.proto[to_utype(PktType::FILE)].nominal_timeout = v.get_uint32();

    else if ( !strcmp(fqn, "stream.ip_cache.idle_timeout") )
        config.flow_cache_cfg.proto[to_utype(PktType::IP)].nominal_timeout = v.get_uint32();

    else if ( !strcmp(fqn, "stream.icmp_cache.idle_timeout") )
        config.flow_cache_cfg.proto[to_utype(PktType::ICMP)].nominal_timeout = v.get_uint32();

    else if ( !strcmp(fqn, "stream.tcp_cache.idle_timeout") )
        config.flow_cache_cfg.proto[to_utype(PktType::TCP)].nominal_timeout = v.get_uint32();

    else if ( !strcmp(fqn, "stream.udp_cache.idle_timeout") )
        config.flow_cache_cfg.proto[to_utype(PktType::UDP)].nominal_timeout = v.get_uint32();

    else if ( v.is("drop_stale_packets") )
    {
        config.drop_stale_packets = v.get_bool();
        if (config.drop_stale_packets)
            c->set_run_flags(RUN_FLAG__DROP_STALE_PACKETS);
        else
            c->clear_run_flags(RUN_FLAG__DROP_STALE_PACKETS);
        return true;
    }
    else
    {
        assert(!strcmp(fqn, "stream.user_cache.idle_timeout"));
        config.flow_cache_cfg.proto[to_utype(PktType::USER)].nominal_timeout = v.get_uint32();
    }

    return true;
}

bool StreamModule::end(const char* fqn, int, SnortConfig* sc)
{
    if ( config.hs_timeout != -1 ) // condition required until stream_tcp.require_3whs is removed
        get_network_parse_policy()->hs_timeout = config.hs_timeout;

    if ( Snort::is_reloading() && strcmp(fqn, MOD_NAME) == 0 )
    {
        StreamReloadResourceManager* reload_resource_manager = new StreamReloadResourceManager;
        if (reload_resource_manager->initialize(config))
            sc->register_reload_handler(reload_resource_manager);
        else
            delete reload_resource_manager;
    }

    return true;
}

void StreamModule::prep_counts(bool)
{ base_prep(); }

void StreamModule::sum_stats(bool dump_stats)
{
    Module::sum_stats(dump_stats);
    if(!dump_stats)
        base_reset();
}

void StreamModule::reset_stats()
{
    base_reset();
    Module::reset_stats();
}

// Stream handler to adjust allocated resources as needed on a config reload
bool StreamReloadResourceManager::initialize(const StreamModuleConfig& config_)
{
    // saving a copy of the config only works here because there is only
    // one stream inspector per packet thread...
#ifdef REG_TEST
    if ( config.footprint != config_.footprint )
    {
        ReloadError("Changing stream.footprint requires a restart.\n");
        return false;
    }
#endif
    config = config_;
    return true;
}

bool StreamReloadResourceManager::tinit()
{
    int max_flows_change =
        config.flow_cache_cfg.max_flows - flow_con->get_flow_cache_config().max_flows;

    if ( max_flows_change )
    {
        if ( max_flows_change < 0 )
            stream_base_stats.reload_total_deletes += abs(max_flows_change);
        else
            stream_base_stats.reload_total_adds += max_flows_change;

        flow_con->set_flow_cache_config(config.flow_cache_cfg);
        return true;
    }

    return false;
}

bool StreamReloadResourceManager::tune_packet_context()
{
    ++stream_base_stats.reload_tuning_packets;
    return tune_resources(max_work);
}

bool StreamReloadResourceManager::tune_idle_context()
{
    ++stream_base_stats.reload_tuning_idle;
    return tune_resources(max_work_idle);
}

bool StreamReloadResourceManager::tune_resources(unsigned work_limit)
{
    // we are done if new max is > currently allocated flow objects
    if ( flow_con->get_flows_allocated() <= config.flow_cache_cfg.max_flows )
        return true;

    unsigned flows_to_delete =
        flow_con->get_flows_allocated() - config.flow_cache_cfg.max_flows;
    if ( flows_to_delete > work_limit )
        flows_to_delete -= flow_con->delete_flows(work_limit);
    else
        flows_to_delete -= flow_con->delete_flows(flows_to_delete);

    return ( flows_to_delete ) ? false : true;
}

void StreamReloadResourceManager::report_progress()
{
    LogMessage("StreamReloadResourceManager[%u] flows allocated %u, max flows %u,"
        " tuning packets %" PRIu64", tuning idle %" PRIu64"\n",
        get_instance_id(), flow_con->get_flows_allocated(), config.flow_cache_cfg.max_flows,
        stream_base_stats.reload_tuning_packets, stream_base_stats.reload_tuning_idle);
}

bool StreamUnloadReloadResourceManager::tinit()
{
    unsigned max_flows = flow_con->get_flow_cache_config().max_flows;
    if (max_flows)
    {
        stream_base_stats.reload_total_deletes += max_flows;
        return true;
    }
    return false;
}

bool StreamUnloadReloadResourceManager::tune_packet_context()
{
    ++stream_base_stats.reload_tuning_packets;
    return tune_resources(max_work);
}

bool StreamUnloadReloadResourceManager::tune_idle_context()
{
    ++stream_base_stats.reload_tuning_idle;
    return tune_resources(max_work_idle);
}

bool StreamUnloadReloadResourceManager::tune_resources(unsigned work_limit)
{
    unsigned flows_to_delete = flow_con->get_flows_allocated();

    if (!flows_to_delete)
        return true;

    if (flows_to_delete > work_limit)
        flows_to_delete -= flow_con->delete_flows(work_limit);
    else
        flows_to_delete -= flow_con->delete_flows(flows_to_delete);

    return (flows_to_delete) ? false : true;
}

void StreamUnloadReloadResourceManager::report_progress()
{
    LogMessage("StreamUnloadReloadResourceManager[%u] flows allocated %u,"
        " tuning packets %" PRIu64", tuning idle %" PRIu64"\n",
        get_instance_id(), flow_con->get_flows_allocated(),
        stream_base_stats.reload_tuning_packets, stream_base_stats.reload_tuning_idle);
}

void StreamModuleConfig::show() const
{
    ConfigLogger::log_value("max_flows", flow_cache_cfg.max_flows);
    ConfigLogger::log_value("max_aux_ip", SnortConfig::get_conf()->max_aux_ip);
    ConfigLogger::log_value("pruning_timeout", flow_cache_cfg.pruning_timeout);
    ConfigLogger::log_value("prune_flows", flow_cache_cfg.prune_flows);
    ConfigLogger::log_limit("require_3whs", hs_timeout, -1, hs_timeout < 0 ? hs_timeout : -1);
    ConfigLogger::log_value("drop_stale_packets", drop_stale_packets ? "enabled" : "disabled");
    ConfigLogger::log_value("held_packet_timeout", hold_time);

    for (int i = to_utype(PktType::IP); i < to_utype(PktType::PDU); ++i)
    {
        std::string tmp;
        tmp += "{ idle_timeout = " + std::to_string(flow_cache_cfg.proto[i].nominal_timeout);
        tmp += " }";

        ConfigLogger::log_value(flow_type_names[i], tmp.c_str());
    }
    {
        std::string tmp;
        tmp += "{ enable = " + (flow_cache_cfg.allowlist_cache ? std::string("true") : std::string("false"));
        tmp += ", move_on_excess = " + (flow_cache_cfg.move_to_allowlist_on_excess ? std::string("true") : std::string("false"));
        tmp += " }";
        ConfigLogger::log_value("allowlist_cache", tmp.c_str());
    }
}

