/* flow2query.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 "flow2query.h"
#include "flow2record.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;

struct Flow2Query::Search_
{
    Flow2Query*                             query;
    std::map<Flow2Query::Key, SummaryIo>&   rows;

    void                                    extract( const Flow2Record* record, int hits, time_t group );
    void                                    find( const Flow2Record* begin, const Flow2Record* end, time_t group );
};

/**************************************************************************************************************/
Flow2Query::Flow2Query() :
    TimeQuery( FLOW2_DIR, 0 ),
    groupByTime_( nullptr ),
    shouldSaveVlans_( false ),
    checkCast_( nullptr ),
    isInverted_( false ),
    isLevel3_( false )
{
}

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

    if ( checkCast_ )
    {
        delete checkCast_;
    }
}

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

    groupByTime_ = new GroupByHour;
    return true;
}

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

    groupByTime_ = new GroupByDay;
    return true;
}

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

    groupByTime_ = new GroupByMonth;
    return true;
}

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

    groupByTime_ = new IgnoreTime;
    return true;
}

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

/**************************************************************************************************************/
bool
Flow2Query::setUnicastFilter()
{
    if ( checkCast_ )
    {
        return false;
    }

    checkCast_ = new CheckUnicast;
    return true;
}

/**************************************************************************************************************/
bool
Flow2Query::setBroadcastFilter()
{
    if ( checkCast_ )
    {
        return false;
    }

    checkCast_ = new CheckBroadcast;
    return true;
}

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

/**************************************************************************************************************/
void
Flow2Query::level3()
{
    isLevel3_ = true;
}

/**************************************************************************************************************/
void
Flow2Query::addMac( const MacAddress& mac )
{
    macs_.insert( mac );
}

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

    if ( qStart >= qEnd || (macs_.empty() && !isInverted_) )
    {
        return;
    }

    readDirectory( {0} );

    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::FLOW2) )
            {
                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 Flow2Record*  records = reinterpret_cast<const Flow2Record*>( file->mapping + file->offset );

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

            search.find( records, nextRecords, regroup(g) );
            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 );
            }
        }
    }
}

/**************************************************************************************************************/
void
Flow2Query::Search_::extract( const Flow2Record* record, int hits, time_t group )
/*
    The hits parameter values interprets as:
        0 - smac is not hit, dmac is not hit
        1 - smac is hit, dmac is not hit
        2 - smac is not hit, dmac is hit
        3 - smac is hit, dmac is hit
*/
{
    int     vlan = ntohs( record->beFlow2.vlan );

    if ( !query->checkCast(record->beFlow2.dmac) || !query->checkVlan(vlan) )
    {
        return;
    }

    Flow2Query::Key     key;

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

    if ( hits & 1 )
    {
        key.host = record->beFlow2.smac;
        key.peer = record->beFlow2.dmac;

        if ( !query->isLevel3() )
        {
            rows[ key ] += SummaryIo( 0, 0, 0, 1, record->summary2.nPackets2(), record->summary2.nBytes2() );
        }
        else if ( record->summary2.nPackets() > 0 )
        {
            rows[ key ] += SummaryIo( 0, 0, 0, 1, record->summary2.nPackets(), record->summary2.nBytes() );
        }
    }

    if ( hits & 2 )
    {
        key.host = record->beFlow2.dmac;
        key.peer = record->beFlow2.smac;

        if ( !query->isLevel3() )
        {
            rows[ key ] += SummaryIo( 1, record->summary2.nPackets2(), record->summary2.nBytes2(), 0, 0, 0 );
        }
        else if ( record->summary2.nPackets() > 0 )
        {
            rows[ key ] += SummaryIo( 1, record->summary2.nPackets(), record->summary2.nBytes(), 0, 0, 0 );
        }
    }
}

/**************************************************************************************************************/
void
Flow2Query::Search_::find( const Flow2Record* begin, const Flow2Record* end, time_t group )
{
    bool    inversion = query->isInverted();

    for ( const Flow2Record* record = begin; record < end; ++record )
    {
        bool    sHit = query->macs().find( record->beFlow2.smac ) != query->macs().end();
        bool    dHit = query->macs().find( record->beFlow2.dmac ) != query->macs().end();
        int     hits = inversion? ( (!sHit && !dHit)? 3 : 0 ) : ( sHit | (dHit << 1) );

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