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

#include "zeek/analyzer/Manager.h"

#include "zeek/Conn.h"
#include "zeek/IntrusivePtr.h"
#include "zeek/RunState.h"
#include "zeek/Val.h"
#include "zeek/packet_analysis/protocol/ip/IPBasedAnalyzer.h"
#include "zeek/packet_analysis/protocol/ip/SessionAdapter.h"

namespace zeek::analyzer {

Manager::ConnIndex::ConnIndex(const IPAddr& _orig, const IPAddr& _resp, uint16_t _resp_p, uint16_t _proto) {
    if ( _orig == IPAddr::v4_unspecified )
        // don't use the IPv4 mapping, use the literal unspecified address
        // to indicate a wildcard
        orig = IPAddr::v6_unspecified;
    else
        orig = _orig;

    resp = _resp;
    resp_p = _resp_p;
    proto = _proto;
}

Manager::ConnIndex::ConnIndex() {
    orig = resp = IPAddr::v4_unspecified;
    resp_p = 0;
    proto = 0;
}

bool Manager::ConnIndex::operator<(const ConnIndex& other) const {
    if ( orig != other.orig )
        return orig < other.orig;

    if ( resp != other.resp )
        return resp < other.resp;

    if ( proto != other.proto )
        return proto < other.proto;

    if ( resp_p != other.resp_p )
        return resp_p < other.resp_p;

    return false;
}

Manager::Manager() : plugin::ComponentManager<analyzer::Component>("Analyzer", "Tag", "AllAnalyzers") {}

Manager::~Manager() {
    // Clean up expected-connection table.
    while ( ! conns_by_timeout.empty() ) {
        ScheduledAnalyzer* a = conns_by_timeout.top();
        conns_by_timeout.pop();
        delete a;
    }
}

void Manager::InitPostScript() {
    const auto& id = detail::global_scope()->Find("PacketAnalyzer::VXLAN::vxlan_ports");

    if ( ! (id && id->GetVal()) )
        reporter->FatalError("PacketAnalyzer::VXLAN::vxlan_ports not defined");

    auto table_val = id->GetVal()->AsTableVal();
    auto port_list = table_val->ToPureListVal();

    for ( auto i = 0; i < port_list->Length(); ++i )
        vxlan_ports.emplace_back(port_list->Idx(i)->AsPortVal()->Port());

    for ( const auto& p : pending_analyzers_for_ports ) {
        if ( ! RegisterAnalyzerForPort(p) )
            reporter->Warning("cannot register analyzer for port %u", std::get<2>(p));
    }

    pending_analyzers_for_ports.clear();

    initialized = true;
}

void Manager::DumpDebug() {
#ifdef DEBUG
    DBG_LOG(DBG_ANALYZER, "Available analyzers after zeek_init():");
    std::list<Component*> all_analyzers = GetComponents();
    for ( std::list<Component*>::const_iterator i = all_analyzers.begin(); i != all_analyzers.end(); ++i )
        DBG_LOG(DBG_ANALYZER, "    %s (%s)", (*i)->Name().c_str(), IsEnabled((*i)->Tag()) ? "enabled" : "disabled");

    DBG_LOG(DBG_ANALYZER, " ");
    DBG_LOG(DBG_ANALYZER, "Analyzers by port:");

    if ( packet_analysis::AnalyzerPtr tcp = packet_mgr->GetAnalyzer("TCP") ) {
        auto* ipba = static_cast<packet_analysis::IP::IPBasedAnalyzer*>(tcp.get());
        ipba->DumpPortDebug();
    }

    if ( packet_analysis::AnalyzerPtr udp = packet_mgr->GetAnalyzer("UDP") ) {
        auto* ipba = static_cast<packet_analysis::IP::IPBasedAnalyzer*>(udp.get());
        ipba->DumpPortDebug();
    }

#endif
}

void Manager::Done() {}

bool Manager::EnableAnalyzer(const zeek::Tag& tag) {
    Component* p = Lookup(tag, false);

    if ( ! p )
        return false;

    DBG_LOG(DBG_ANALYZER, "Enabling analyzer %s", p->Name().c_str());
    p->SetEnabled(true);

    return true;
}

bool Manager::EnableAnalyzer(EnumVal* val) {
    Component* p = Lookup(val, false);

    if ( ! p )
        return false;

    DBG_LOG(DBG_ANALYZER, "Enabling analyzer %s", p->Name().c_str());
    p->SetEnabled(true);

    return true;
}

bool Manager::DisableAnalyzer(const zeek::Tag& tag) {
    Component* p = Lookup(tag, false);

    if ( ! p )
        return false;

    DBG_LOG(DBG_ANALYZER, "Disabling analyzer %s", p->Name().c_str());
    p->SetEnabled(false);

    return true;
}

bool Manager::DisableAnalyzer(EnumVal* val) {
    Component* p = Lookup(val);

    if ( ! p )
        return false;

    DBG_LOG(DBG_ANALYZER, "Disabling analyzer %s", p->Name().c_str());
    p->SetEnabled(false);

    return true;
}

void Manager::DisableAllAnalyzers() {
    DBG_LOG(DBG_ANALYZER, "Disabling all analyzers");

    std::list<Component*> all_analyzers = GetComponents();
    for ( std::list<Component*>::const_iterator i = all_analyzers.begin(); i != all_analyzers.end(); ++i )
        (*i)->SetEnabled(false);
}

zeek::Tag Manager::GetAnalyzerTag(const char* name) { return GetComponentTag(name); }

bool Manager::IsEnabled(const zeek::Tag& tag) {
    if ( ! tag )
        return false;

    Component* p = Lookup(tag);

    if ( ! p )
        return false;

    return p->Enabled();
}

bool Manager::IsEnabled(EnumVal* val) {
    Component* p = Lookup(val);

    if ( ! p )
        return false;

    return p->Enabled();
}

bool Manager::RegisterAnalyzerForPort(EnumVal* val, PortVal* port) {
    Component* p = Lookup(val);

    if ( ! p )
        return false;

    return RegisterAnalyzerForPort(p->Tag(), port->PortType(), port->Port());
}

bool Manager::UnregisterAnalyzerForPort(EnumVal* val, PortVal* port) {
    Component* p = Lookup(val);

    if ( ! p )
        return false;

    return UnregisterAnalyzerForPort(p->Tag(), port->PortType(), port->Port());
}

bool Manager::RegisterAnalyzerForPort(const zeek::Tag& tag, TransportProto proto, uint32_t port) {
    if ( initialized )
        return RegisterAnalyzerForPort(std::make_tuple(tag, proto, port));
    else {
        // Cannot register these before PostScriptInit() has run because we
        // depend on packet analysis having been set up. That also means we don't have
        // a reliable return value, for now we just assume it's working.
        pending_analyzers_for_ports.emplace(tag, proto, port);
        return true;
    }
}

bool Manager::RegisterAnalyzerForPort(const std::tuple<zeek::Tag, TransportProto, uint32_t>& p) {
    const auto& [tag, proto, port] = p;

    // TODO: this class is becoming more generic and removing a lot of the
    // checks for protocols, but this part might need to stay like this.
    packet_analysis::AnalyzerPtr analyzer;
    if ( proto == TRANSPORT_TCP )
        analyzer = packet_mgr->GetAnalyzer("TCP");
    else if ( proto == TRANSPORT_UDP )
        analyzer = packet_mgr->GetAnalyzer("UDP");

    if ( ! analyzer )
        return false;

    auto* ipba = static_cast<packet_analysis::IP::IPBasedAnalyzer*>(analyzer.get());

    return ipba->RegisterAnalyzerForPort(tag, port);
}

bool Manager::UnregisterAnalyzerForPort(const zeek::Tag& tag, TransportProto proto, uint32_t port) {
    if ( auto i = pending_analyzers_for_ports.find(std::make_tuple(tag, proto, port));
         i != pending_analyzers_for_ports.end() )
        pending_analyzers_for_ports.erase(i);

    // TODO: this class is becoming more generic and removing a lot of the
    // checks for protocols, but this part might need to stay like this.
    packet_analysis::AnalyzerPtr analyzer;
    if ( proto == TRANSPORT_TCP )
        analyzer = packet_mgr->GetAnalyzer("TCP");
    else if ( proto == TRANSPORT_UDP )
        analyzer = packet_mgr->GetAnalyzer("UDP");

    if ( ! analyzer )
        return false;

    auto* ipba = static_cast<packet_analysis::IP::IPBasedAnalyzer*>(analyzer.get());

    return ipba->UnregisterAnalyzerForPort(tag, port);
}

Analyzer* Manager::InstantiateAnalyzer(const zeek::Tag& tag, Connection* conn) {
    Component* c = Lookup(tag);

    if ( ! c ) {
        reporter->InternalWarning("request to instantiate unknown analyzer");
        return nullptr;
    }

    if ( ! c->Enabled() )
        return nullptr;

    if ( ! c->Factory() ) {
        reporter->InternalWarning("analyzer %s cannot be instantiated dynamically", GetComponentName(tag).c_str());
        return nullptr;
    }

    Analyzer* a = c->Factory()(conn);

    if ( ! a ) {
        reporter->InternalWarning("analyzer instantiation failed");
        return nullptr;
    }

    a->SetAnalyzerTag(tag);

    return a;
}

Analyzer* Manager::InstantiateAnalyzer(const char* name, Connection* conn) {
    zeek::Tag tag = GetComponentTag(name);
    return tag ? InstantiateAnalyzer(tag, conn) : nullptr;
}

bool Manager::AnalyzerTransportMatches(const zeek::Tag& tag, TransportProto transport) const {
    Component* c = Lookup(tag);
    return ! c || c->Transport() == TRANSPORT_UNKNOWN || c->Transport() == transport;
}

void Manager::ExpireScheduledAnalyzers() {
    if ( ! run_state::network_time )
        return;

    while ( ! conns_by_timeout.empty() ) {
        ScheduledAnalyzer* a = conns_by_timeout.top();

        if ( a->queue_timeout > run_state::network_time )
            return;

        conns_by_timeout.pop();

        // If we scheduled the same analyzer more than once, the *real* timeout might not
        // have expired. Push it back if so.
        assert(a->queue_timeout <= a->timeout);
        if ( a->timeout > run_state::network_time ) {
            a->queue_timeout = a->timeout;
            conns_by_timeout.push(a);
            continue;
        }

        std::pair<conns_map::iterator, conns_map::iterator> all = conns.equal_range(a->conn);

        bool found = false;

        for ( conns_map::iterator i = all.first; i != all.second; i++ ) {
            if ( i->second != a )
                continue;

            conns.erase(i);

            DBG_LOG(DBG_ANALYZER, "Expiring expected analyzer %s for connection %s",
                    analyzer_mgr->GetComponentName(a->analyzer).c_str(),
                    fmt_conn_id(a->conn.orig, 0, a->conn.resp, a->conn.resp_p));

            delete a;
            found = true;
            break;
        }

        assert(found);
    }
}

void Manager::ScheduleAnalyzer(const IPAddr& orig, const IPAddr& resp, uint16_t resp_p, TransportProto proto,
                               const zeek::Tag& analyzer, double timeout) {
    if ( ! run_state::network_time ) {
        reporter->Warning("cannot schedule analyzers before processing begins; ignored");
        return;
    }

    assert(timeout);

    // Use the chance to see if the oldest entry is already expired.
    ExpireScheduledAnalyzers();

    ConnIndex conn(orig, resp, resp_p, proto);
    const double expiration = run_state::network_time + timeout;
    const auto scheduled = conns.equal_range(conn);

    // Find any previously scheduled analyzer entry for the connection with a matching
    // tag and update its timeout in place. This does not change its queue position.
    // Instead, the analyzer will get re-queued on expiration, if necessary.
    for ( auto it = scheduled.first; it != scheduled.second; ++it ) {
        ScheduledAnalyzer* a = it->second;

        if ( a->analyzer != analyzer )
            continue;

        // It's the same, so adjust the real timeout and return.
        // This does NOT change the queue position!
        if ( a->timeout < expiration )
            a->timeout = expiration;

        return;
    }

    ScheduledAnalyzer* a = new ScheduledAnalyzer;
    a->conn = conn;
    a->analyzer = analyzer;
    a->timeout = expiration;
    a->queue_timeout = expiration;

    conns.insert(std::make_pair(a->conn, a));
    conns_by_timeout.push(a);
}

void Manager::ScheduleAnalyzer(const IPAddr& orig, const IPAddr& resp, uint16_t resp_p, TransportProto proto,
                               const char* analyzer, double timeout) {
    zeek::Tag tag = GetComponentTag(analyzer);

    if ( tag != zeek::Tag() )
        ScheduleAnalyzer(orig, resp, resp_p, proto, tag, timeout);
}

void Manager::ScheduleAnalyzer(const IPAddr& orig, const IPAddr& resp, PortVal* resp_p, Val* analyzer, double timeout) {
    EnumValPtr ev{NewRef{}, analyzer->AsEnumVal()};
    return ScheduleAnalyzer(orig, resp, resp_p->Port(), resp_p->PortType(), zeek::Tag(std::move(ev)), timeout);
}

Manager::tag_set Manager::GetScheduled(const Connection* conn) {
    ConnIndex c(conn->OrigAddr(), conn->RespAddr(), ntohs(conn->RespPort()), conn->ConnTransport());

    std::pair<conns_map::iterator, conns_map::iterator> all = conns.equal_range(c);

    tag_set result;

    for ( conns_map::iterator i = all.first; i != all.second; i++ )
        result.insert(i->second->analyzer);

    // Try wildcard for originator.
    c.orig = IPAddr::v6_unspecified;
    all = conns.equal_range(c);

    for ( conns_map::iterator i = all.first; i != all.second; i++ ) {
        if ( i->second->timeout > run_state::network_time )
            result.insert(i->second->analyzer);
    }

    // We don't delete scheduled analyzers here. They will be expired
    // eventually.
    return result;
}

bool Manager::ApplyScheduledAnalyzers(Connection* conn, bool init, packet_analysis::IP::SessionAdapter* parent) {
    if ( ! parent )
        parent = conn->GetSessionAdapter();

    if ( ! parent )
        return false;

    tag_set expected = GetScheduled(conn);

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

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

        if ( ! analyzer )
            continue;

        parent->AddChildAnalyzer(analyzer, init);

        if ( scheduled_analyzer_applied )
            conn->EnqueueEvent(scheduled_analyzer_applied, nullptr, conn->GetVal(), tag.AsVal());

        DBG_ANALYZER_ARGS(conn, "activated %s analyzer as scheduled", analyzer_mgr->GetComponentName(tag).c_str());
    }

    return ! expected.empty();
}

} // namespace zeek::analyzer
