/* flowquery.cc
****************************************************************************************************************
****************************************************************************************************************

    Copyright (C) 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2020, 2021 Askar
    Almatov

    This file is part of Trafficmeter.

    Trafficmeter is free software: you can redistribute it and/or modify it under the terms of the GNU General
    Public License as published by the Free Software Foundation, either version 3 of the License, or (at your
    option) any later version.

    Trafficmeter 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 Trafficmeter. If not, see
    <http://www.gnu.org/licenses/>.
*/

#include <cassert>
#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/mman.h>

#include "fileheader.h"
#include "flowquery.h"
#include "flowrecord.h"
#include "tmdefs.h"

using namespace iclass;
using namespace trafficmeter;

using std::map;
using std::max;
using std::min;
using std::runtime_error;
using std::string;

static const int    SERVICES_[] =
{
    53,                 // dns
    80,                 // http
    443,                // https
    110, 995,           // pop
    143, 993,           // imap
    25, 465,            // smtp
    23399,              // skype
    5060,               // sip
    4569,               // iax2
    1719, 1720,         // h.323
    22,                 // ssh
    23,                 // telnet
    5190,               // icq
    5050,               // yahoo-im
    2041,               // mail.ru-im
    1863,               // msn-im
    194, 994,           // irc
    20, 21, 989, 990,   // ftp
    119, 563,           // nntp
    123,                // ntp
    2427,               // mgcp
    2000,               // skinny
    2944,               // h.248
    1293,               // ipsec
    1812,               // radius
    49,                 // tacacs
    3868,               // diameter
    137, 138, 139,      // netbios
    3128, 8080,         // webcache
    3306,               // mysql
    5432,               // postgresql
    1433,               // mssql
    1047,               // neod1-trojan
    0
};

struct FlowQuery::Search
{
    FlowQuery*                              query;
    std::map<FlowQuery::Key, SummaryIo>*    rows;

    virtual         ~Search() {}
    void            extract( const FlowRecord* record, int hits, time_t group );
    virtual void    find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash ) = 0;
};

struct SearchByIp_ : public FlowQuery::Search
{
    uint32_t        beIp;

                    SearchByIp_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows );
    virtual         ~SearchByIp_() {}
    virtual void    find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash ) override;
};

struct SearchByNetwork_ : public FlowQuery::Search
{
    uint32_t        beNet;
    uint32_t        beMask;

                    SearchByNetwork_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows );
    virtual         ~SearchByNetwork_() {}
    virtual void    find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash ) override;
};

struct SearchBySet_ : public FlowQuery::Search
{
    iplookup4t      lookup;

                    SearchBySet_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows );
    virtual         ~SearchBySet_() {}
    virtual void    find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash ) override;
};

/**************************************************************************************************************/
FlowQuery::FlowQuery() :
    TimeQuery( FLOW_DIR, SUFFIX_HASH_LENGTH ),
    groupByTime_( nullptr ),
    shouldSaveVlans_( false ),
    shouldSaveAllPorts_( false ),
    isInverted_( false )
{
}

/**************************************************************************************************************/
FlowQuery::~FlowQuery()
{
    if ( groupByTime_ )
    {
        delete groupByTime_;
    }
}

/**************************************************************************************************************/
bool
FlowQuery::timeByHour()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByHour;
    return true;
}

/**************************************************************************************************************/
bool
FlowQuery::timeByDay()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByDay;
    return true;
}

/**************************************************************************************************************/
bool
FlowQuery::timeByMonth()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByMonth;
    return true;
}

/**************************************************************************************************************/
bool
FlowQuery::hideTime()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new IgnoreTime;
    return true;
}

/**************************************************************************************************************/
void
FlowQuery::saveVlans()
{
    shouldSaveVlans_ = true;
}

/**************************************************************************************************************/
void
FlowQuery::saveAllPorts()
{
    shouldSaveAllPorts_ = true;
}

/**************************************************************************************************************/
void
FlowQuery::invert()
{
    isInverted_ = true;

    if ( hashes_.size() < (1U<<SUFFIX_HASH_LENGTH) )
    {
        int     hash = 0;
        int     nHashes = 1U << SUFFIX_HASH_LENGTH;

        while ( nHashes-- )
        {
            hashes_.insert( hash++ );
        }
    }
}

/**************************************************************************************************************/
void
FlowQuery::addNetwork( ipnet4t network )
{
    networks_.insert( network );

    if ( hashes_.size() < (1U<<SUFFIX_HASH_LENGTH) )
    {
        int     hash = network.front() & ( (1U << SUFFIX_HASH_LENGTH) - 1 );
        int     nHashes = 1U << min( network.nodebits(), SUFFIX_HASH_LENGTH );

        while ( nHashes-- )
        {
            hashes_.insert( hash++ );
        }
    }
}

/**************************************************************************************************************/
void
FlowQuery::fetch( map<Key,SummaryIo>& summaries )
{
    time_t      qStart = begin().unixtime();
    time_t      qEnd = end().unixtime();
    Search*     search;

    if ( qStart >= qEnd || networks_.empty() )
    {
        return;
    }

    if ( networks_.size() > 1 || isInverted_ )
    {
        search = new SearchBySet_( this, &summaries );
    }
    else if ( networks_.begin()->nodebits() )
    {
        search = new SearchByNetwork_( this, &summaries );
    }
    else
    {
        search = new SearchByIp_( this, &summaries );
    }

    readDirectory( hashes_ );

    TimeQuery::Range    files = getFiles( qEnd );

    for ( auto file = files.first; file != files.second; ++file )
    {
        if ( file->isScanned )
        {
            continue;
        }

        FileHeader  header;
        time_t      fStart;
        time_t      fEnd;
        long        mapOffset;

        if ( file->fd < 0 )
        {
            Path&&  path( filePath(*file) );

            file->fd = open( path.c_str(), O_RDONLY );

            if ( file->fd == -1 || !header.read(file->fd) )
            {
                throw runtime_error( "file read error: " + path + " - " + strerror(errno) );
            }

            if ( !header.isCompatible(FileHeader::Type::FLOW) )
            {
                throw runtime_error( "incompatible file format: " + path );
            }

            fStart = max( moduloRoundDown(qStart, GROUP_FACTOR), header.oldestGroup() );
            fEnd = min( moduloRoundDown(qEnd-1, GROUP_FACTOR), header.newestGroup() ) + GROUP_FACTOR;
            mapOffset = header.groupOffset( fStart );
            file->mapSize = header.groupOffset( fEnd ) - mapOffset;

            if ( file->mapSize <= 0 )
            {
                file->isScanned = true;
                continue;
            }

            file->offset = mapOffset % getpagesize();
            mapOffset -= file->offset;
            file->mapSize += file->offset;
            file->mapping = static_cast<uint8_t*>( mmap(nullptr, file->mapSize, PROT_READ, MAP_SHARED, file->fd, mapOffset) );

            if ( file->mapping == MAP_FAILED )
            {
                throw runtime_error( "mmap error: " + path + " - " + strerror(errno) );
            }
        }

        const FlowRecord*   records = reinterpret_cast<const FlowRecord*>( file->mapping + file->offset );

        for ( time_t g = fStart; g < fEnd; g += GROUP_FACTOR )
        {
            long                nextOffset = header.groupOffset( g + GROUP_FACTOR ) - mapOffset;
            const FlowRecord*   nextRecords = reinterpret_cast<const FlowRecord*>( file->mapping + nextOffset );

            search->find( records, nextRecords, regroup(g), file->hash() );
            records = nextRecords;
        }

        file->isScanned = true;
    }

    for ( auto file = files.first; file != files.second; ++file )
    {
        if ( file->isScanned )
        {
            if ( file->mapping != nullptr )
            {
                munmap( file->mapping, file->mapSize );
            }

            if ( file->fd >= 0 )
            {
                close( file->fd );
            }
        }
    }

    delete search;
}

/**************************************************************************************************************/
void
FlowQuery::Search::extract( const FlowRecord* record, int hits, time_t group )
/*
    The hits parameter values interprets as:
        0 - sip is not hit, dip is not hit
        1 - sip is hit, dip is not hit
        2 - sip is not hit, dip is hit
        3 - sip is hit, dip is hit
*/
{
    int     vlan = ntohs( record->beFlow.vlan );

    if ( !query->checkVlan(vlan) )
    {
        return;
    }

    FlowQuery::Key  key;

    key.group = group;
    key.vlan = query->shouldSaveVlans()? vlan : 0;
    key.proto = ntohs( record->beFlow.proto );

    int     sport = ntohs( record->beFlow.sport );
    int     dport = ntohs( record->beFlow.dport );

    if ( key.proto != 6 && key.proto != 17 )
    {
        sport = -1;
        dport = -1;
    }
    else if ( !query->shouldSaveAllPorts() )
    {
        bool    shouldHideSport( true );
        bool    shouldHideDport( true );

        for ( const int* service = SERVICES_; *service != 0; ++service )
        {
            if ( sport == *service )
            {
                shouldHideSport = false;
            }

            if ( dport == *service )
            {
                shouldHideDport = false;
            }
        }

        if ( shouldHideSport )
        {
            sport = -1;
        }

        if ( shouldHideDport )
        {
            dport = -1;
        }
    }

    ipaddr4t    sip( ntohl(record->beFlow.sip) );
    ipaddr4t    dip( ntohl(record->beFlow.dip) );

    if ( hits & 1 )
    {
        key.host = sip;
        key.peer = dip;
        key.hport = sport;
        key.pport = dport;
        (*rows)[ key ] += SummaryIo( 0, 0, 0, 1, record->summary.nPackets(), record->summary.nBytes() );
    }

    if ( hits & 2 )
    {
        key.host = dip;
        key.peer = sip;
        key.hport = dport;
        key.pport = sport;
        (*rows)[ key ] += SummaryIo( 1, record->summary.nPackets(), record->summary.nBytes(), 0, 0, 0 );
    }
}

/**************************************************************************************************************/
SearchByIp_::SearchByIp_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows )
{
    Search::query = query;
    Search::rows = rows;
    beIp = htonl( query->networks().begin()->front() );
}

/**************************************************************************************************************/
void
SearchByIp_::find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash )
{
    for ( const FlowRecord* record = begin; record < end; ++record )
    {
        int     hits = ( record->beFlow.sip == beIp ) | ( (record->beFlow.dip == beIp) << 1 );

        if ( hits )
        {
            extract( record, hits, group );
        }
    }
}

/**************************************************************************************************************/
SearchByNetwork_::SearchByNetwork_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows )
{
    Search::query = query;
    Search::rows = rows;
    beNet = htonl( query->networks().begin()->front() );
    beMask = htonl( query->networks().begin()->mask() );
}

/**************************************************************************************************************/
void
SearchByNetwork_::find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash )
{
    uint32_t    beHash = htonl( hash );
    uint32_t    beHashMask = htonl( (1U << SUFFIX_HASH_LENGTH) - 1 );

    for ( const FlowRecord* record = begin; record < end; ++record )
    {
        bool    sHit = ( record->beFlow.sip & beMask ) == beNet && ( record->beFlow.sip & beHashMask ) == beHash;
        bool    dHit = ( record->beFlow.dip & beMask ) == beNet && ( record->beFlow.dip & beHashMask ) == beHash;
        int     hits = sHit | ( dHit << 1 );

        if ( hits )
        {
            extract( record, hits, group );
        }
    }
}

/**************************************************************************************************************/
SearchBySet_::SearchBySet_( FlowQuery* query, std::map<FlowQuery::Key, SummaryIo>* rows ) :
    lookup( query->networks() )
{
    Search::query = query;
    Search::rows = rows;
}

/**************************************************************************************************************/
void
SearchBySet_::find( const FlowRecord* begin, const FlowRecord* end, time_t group, int hash )
{
    uint32_t    beHash = htonl( hash );
    uint32_t    beHashMask = htonl( (1U << SUFFIX_HASH_LENGTH) - 1 );

    for ( const FlowRecord* record = begin; record < end; ++record )
    {
        bool  sHit = lookup.contains( ntohl(record->beFlow.sip) );
        bool  dHit = lookup.contains( ntohl(record->beFlow.dip) );
        int   hits = query->isInverted()? ( (!sHit && !dHit)? 3 : 0 ) : ( sHit | (dHit << 1) );

        hits &= ( (record->beFlow.sip & beHashMask) == beHash ) | ( ((record->beFlow.dip&beHashMask) == beHash) << 1 );

        if ( hits )
        {
            extract( record, hits, group );
        }
    }
}
