// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/packet_analysis/protocol/ip/IPBasedAnalyzer.h"

#include <cinttypes>

#include "zeek/Conn.h"
#include "zeek/RunState.h"
#include "zeek/Val.h"
#include "zeek/analyzer/Manager.h"
#include "zeek/analyzer/protocol/pia/PIA.h"
#include "zeek/conn_key/Manager.h"
#include "zeek/packet_analysis/protocol/ip/SessionAdapter.h"
#include "zeek/packet_analysis/protocol/ip/conn_key/IPBasedConnKey.h"
#include "zeek/plugin/Manager.h"
#include "zeek/session/Manager.h"

using namespace zeek;
using namespace zeek::packet_analysis::IP;

IPBasedAnalyzer::IPBasedAnalyzer(const char* name, TransportProto proto, uint32_t mask, bool report_unknown_protocols)
    : zeek::packet_analysis::Analyzer(name, report_unknown_protocols), transport(proto), server_port_mask(mask) {}

IPBasedAnalyzer::~IPBasedAnalyzer() {
    for ( const auto& mapping : analyzers_by_port )
        delete mapping.second;
}

bool IPBasedAnalyzer::AnalyzePacket(size_t len, const uint8_t* data, Packet* pkt) {
    static IPBasedConnKeyPtr key; // Note, this is static for reuse:
    if ( ! key ) {
        ConnKeyPtr ck = conn_key_mgr->GetFactory().NewConnKey();

        // The IPBasedAnalyzer requires a factory that produces IPBasedConnKey instances.
        // We could check with dynamic_cast, but that's probably slow, so assume plugin
        // providers know what they're doing here and anyhow, we don't really have analyzers
        // that instantiate non-IP connections today and definitely not here!
        key = IPBasedConnKeyPtr(static_cast<IPBasedConnKey*>(ck.release()));
    }

    if ( ! InitConnKey(len, data, pkt, *key) )
        return false;

    key->Init(*pkt);

    const std::shared_ptr<IP_Hdr>& ip_hdr = pkt->ip_hdr;
    auto src_addr = key->SrcAddr();
    auto src_port = key->SrcPort();

    Connection* conn = session_mgr->FindConnection(*key);

    if ( ! conn ) {
        conn = NewConn(std::move(key), pkt);
        if ( conn )
            session_mgr->Insert(conn, false);
    }
    else {
        if ( conn->IsReuse(run_state::processing_start_time, ip_hdr->Payload()) ) {
            conn->Event(connection_reused, nullptr);

            session_mgr->Remove(conn);
            conn = NewConn(std::move(key), pkt);
            if ( conn )
                session_mgr->Insert(conn, false);
        }
        else {
            conn->CheckEncapsulation(pkt->encap);
        }
    }

    if ( ! conn )
        return false;

    // If we successfully made a connection for this packet that means it'll eventually
    // get logged, which means we can mark this packet as having been processed.
    pkt->processed = true;

    bool is_orig = (src_addr == conn->OrigAddr()) && (src_port == conn->OrigPort());
    pkt->is_orig = is_orig;

    conn->CheckFlowLabel(is_orig, ip_hdr->FlowLabel());

    zeek::ValPtr pkt_hdr_val;

    if ( ipv6_ext_headers && ip_hdr->NumHeaders() > 1 ) {
        pkt_hdr_val = ip_hdr->ToPktHdrVal();
        conn->EnqueueEvent(ipv6_ext_headers, nullptr, conn->GetVal(), pkt_hdr_val);
    }

    if ( new_packet )
        conn->EnqueueEvent(new_packet, nullptr, conn->GetVal(),
                           pkt_hdr_val ? std::move(pkt_hdr_val) : ip_hdr->ToPktHdrVal());

    conn->SetRecordPackets(true);
    conn->SetRecordContents(true);

    const u_char* payload = pkt->ip_hdr->Payload();

    run_state::current_timestamp = run_state::processing_start_time;
    run_state::current_pkt = pkt;

    const auto* adapter = conn->GetSessionAdapter();
    if ( adapter->Skipping() ) {
        adapter->TapPacket(pkt, PacketAction::Skip, SkipReason::SkipProcessing);
        return true;
    }

    DeliverPacket(conn, run_state::processing_start_time, is_orig, len, pkt);

    run_state::current_timestamp = 0;
    run_state::current_pkt = nullptr;

    // If the packet is reassembled, disable packet dumping because the
    // pointer math to dump the data wouldn't work.
    if ( pkt->ip_hdr->Reassembled() )
        pkt->dump_packet = false;
    else if ( conn->RecordPackets() ) {
        pkt->dump_packet = true;

        // If we don't want the content, set the dump size to include just
        // the header.
        if ( ! conn->RecordContents() )
            pkt->dump_size = payload - pkt->data;
    }

    return true;
}

bool IPBasedAnalyzer::CheckHeaderTrunc(size_t min_hdr_len, size_t remaining, Packet* packet) {
    // If segment offloading or similar is enabled, the payload len will return 0.
    // Thus, let's ignore that case.
    if ( packet->ip_hdr->PayloadLen() && packet->ip_hdr->PayloadLen() < min_hdr_len ) {
        Weird("truncated_header", packet);
        return false;
    }
    else if ( remaining < min_hdr_len ) {
        Weird("internally_truncated_header", packet);
        return false;
    }

    return true;
}

bool IPBasedAnalyzer::IsLikelyServerPort(uint32_t port) const {
    // We keep a cached in-core version of the table to speed up the lookup.
    static std::set<zeek_uint_t> port_cache;
    static bool have_cache = false;

    if ( ! have_cache ) {
        auto likely_server_ports = id::find_val<TableVal>("likely_server_ports");
        auto lv = likely_server_ports->ToPureListVal();
        for ( int i = 0; i < lv->Length(); i++ )
            port_cache.insert(lv->Idx(i)->InternalUnsigned());
        have_cache = true;
    }

    // We exploit our knowledge of PortVal's internal storage mechanism here.
    port |= server_port_mask;

    return port_cache.contains(port);
}

zeek::Connection* IPBasedAnalyzer::NewConn(IPBasedConnKeyPtr key, const Packet* pkt) {
    auto src_p = ntohs(key->SrcPort());
    auto dst_p = ntohs(key->DstPort());
    bool flip = false;

    if ( ! WantConnection(src_p, dst_p, pkt->ip_hdr->Payload(), flip) )
        return nullptr;

    Connection* conn = new Connection(std::move(key), run_state::processing_start_time, pkt->ip_hdr->FlowLabel(), pkt);
    conn->SetTransport(transport);

    if ( flip && ! conn->RespAddr().IsBroadcast() )
        conn->FlipRoles();

    BuildSessionAnalyzerTree(conn);

    if ( new_connection )
        conn->Event(new_connection, nullptr);

    return conn;
}

void IPBasedAnalyzer::BuildSessionAnalyzerTree(Connection* conn) {
    SessionAdapter* root = MakeSessionAdapter(conn);
    analyzer::pia::PIA* pia = MakePIA(conn);

    bool scheduled = analyzer_mgr->ApplyScheduledAnalyzers(conn, false, root);

    // Hmm... Do we want *just* the expected analyzer, or all
    // other potential analyzers as well?  For now we only take
    // the scheduled ones.
    if ( ! scheduled ) { // Let's see if it's a port we know.
        if ( ! analyzers_by_port.empty() && ! zeek::detail::dpd_ignore_ports ) {
            int resp_port = ntohs(conn->RespPort());
            std::set<zeek::Tag>* ports = LookupPort(resp_port, false);

            if ( ports ) {
                for ( const auto& port : *ports ) {
                    if ( ! analyzer_mgr->AnalyzerTransportMatches(port, conn->ConnTransport()) ) {
                        reporter->Error("refusing to attach analyzer %s to %s connection with incompatible transport",
                                        analyzer_mgr->GetComponentName(port).c_str(),
                                        transport_proto_string(conn->ConnTransport()));
                        continue;
                    }

                    analyzer::Analyzer* analyzer = analyzer_mgr->InstantiateAnalyzer(port, conn);

                    if ( ! analyzer )
                        continue;

                    root->AddChildAnalyzer(analyzer, false);
                    DBG_ANALYZER_ARGS(conn, "activated %s analyzer due to port %d",
                                      analyzer_mgr->GetComponentName(port).c_str(), resp_port);
                }
            }
        }
    }

    // Make analyzers added above through known ports visible via GetChildren()
    root->AppendNewChildren();

    root->AddExtraAnalyzers(conn);

    if ( pia )
        root->AddChildAnalyzer(pia->AsAnalyzer());

    conn->SetSessionAdapter(root, pia);
    root->Init();
    root->InitChildren();

    PLUGIN_HOOK_VOID(HOOK_SETUP_ANALYZER_TREE, HookSetupAnalyzerTree(conn));
}

bool IPBasedAnalyzer::RegisterAnalyzerForPort(const zeek::Tag& tag, uint32_t port) {
    tag_set* l = LookupPort(port, true);

    if ( ! l )
        return false;

#ifdef DEBUG
    const char* name = analyzer_mgr->GetComponentName(tag).c_str();
    DBG_LOG(DBG_ANALYZER, "Registering analyzer %s for port %" PRIu32 "/%d", name, port, transport);
#endif

    l->insert(tag);
    return true;
}

bool IPBasedAnalyzer::UnregisterAnalyzerForPort(const zeek::Tag& tag, uint32_t port) {
    tag_set* l = LookupPort(port, true);

    if ( ! l )
        return true; // still a "successful" unregistration

#ifdef DEBUG
    const char* name = analyzer_mgr->GetComponentName(tag).c_str();
    DBG_LOG(DBG_ANALYZER, "Unregistering analyzer %s for port %" PRIu32 "/%d", name, port, transport);
#endif

    l->erase(tag);
    return true;
}

IPBasedAnalyzer::tag_set* IPBasedAnalyzer::LookupPort(uint32_t port, bool add_if_not_found) {
    analyzer_map_by_port::const_iterator i = analyzers_by_port.find(port);

    if ( i != analyzers_by_port.end() )
        return i->second;

    if ( ! add_if_not_found )
        return nullptr;

    tag_set* l = new tag_set{};
    analyzers_by_port.insert(std::make_pair(port, l));
    return l;
}

void IPBasedAnalyzer::DumpPortDebug() {
    for ( const auto& mapping : analyzers_by_port ) {
        std::string s;

        for ( const auto& tag : *(mapping.second) )
            s += std::string(analyzer_mgr->GetComponentName(tag)) + " ";

        DBG_LOG(DBG_ANALYZER, "    %d/%s: %s", mapping.first, transport_proto_string(transport), s.c_str());
    }
}

TableValPtr IPBasedAnalyzer::ignore_checksums_nets_table = nullptr;

void IPBasedAnalyzer::SetIgnoreChecksumsNets(TableValPtr t) {
    IPBasedAnalyzer::ignore_checksums_nets_table = std::move(t);
}

TableValPtr IPBasedAnalyzer::GetIgnoreChecksumsNets() {
    if ( ! IPBasedAnalyzer::ignore_checksums_nets_table )
        IPBasedAnalyzer::ignore_checksums_nets_table = zeek::id::find_val<TableVal>("ignore_checksums_nets");
    return IPBasedAnalyzer::ignore_checksums_nets_table;
}
