/* packetquery.cc
****************************************************************************************************************
****************************************************************************************************************

    Copyright (C) 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 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 <algorithm>
#include <cstring>
#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/resource.h>

#include "path.h"
#include "packetquery.h"
#include "tmdefs.h"

using namespace iclass;
using namespace trafficmeter;

using std::pair;
using std::runtime_error;
using std::string;
using std::vector;

/**************************************************************************************************************/
PacketQuery::PacketQuery() :
    TimeQuery( PACKET_DIR, SUFFIX_HASH_LENGTH ),
    hash_( -1 )
{
}

/**************************************************************************************************************/
void
PacketQuery::fetch( PcapWriting& pcap, Stat& statistics )
{
    time_t  qStart = begin().unixtime();
    time_t  qEnd = end().unixtime();

    if ( qStart >= qEnd || hash_ < 0 )
    {
        return;
    }

    readDirectory( {hash_} );

    for ( time_t timeBound = qStart; ++timeBound <= qEnd; )
    {
        vector<pair<Timestamp,const PacketRecord*>>     packets;
        TimeQuery::Range                                files = getFiles( timeBound );

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

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

                if ( (file->fd = open(path.c_str(),O_RDONLY)) == -1 )
                {
                    throw runtime_error( "file read error: " + path + " - " + strerror(errno) );
                }

                file->mapSize = path.fileSize();
                file->mapping = static_cast<uint8_t*>( mmap(nullptr, file->mapSize, PROT_READ, MAP_SHARED, file->fd, 0) );

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

            Timestamp               searchStart( qStart, 0 );
            Timestamp               searchEnd( timeBound, 0 );
            const PacketRecord*     records = reinterpret_cast<const PacketRecord*>( file->mapping + file->offset );
            const PacketRecord*     recordsEnd = reinterpret_cast<const PacketRecord*>( file->mapping + file->mapSize );
            const PacketRecord*     record;

            for ( record = records; record < recordsEnd; record = record->next() )
            {
                if ( record->timestamp < searchStart )
                {
                    continue;
                }
                else if ( record->timestamp >= searchEnd )
                {
                    break;
                }

                Metadata    meta( Timestamp(record->timestamp), record->data2(), record->nBytes2 );
                bool        sHit( meta.sip == ip_ );
                bool        dHit( meta.dip == ip_ );

                if ( (sHit || dHit) && checkVlan(meta.vlan) )
                {
                    packets.push_back( {Timestamp(meta.timestamp), record} );

                    if ( !pcap.isConsole() )
                    {
                        statistics.add( meta, sHit, dHit );
                    }
                }
            }

            file->offset = reinterpret_cast<const uint8_t*>( record ) - file->mapping;

            if ( record >= recordsEnd || record->timestamp > Timestamp(qEnd, 0) )
            {
                file->isScanned = true;
            }
        }

        sort( packets.begin(), packets.end() );

        for ( auto p : packets )
        {
            pcap.writePacket( p.first, p.second->data2(), p.second->nBytes2 );
        }

        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
PacketQuery::setIp( ipaddr4t Ip )
{
    ip_ = Ip;
    hash_ = ip_ & ( (1U << SUFFIX_HASH_LENGTH) - 1 );
}

/**************************************************************************************************************/
iclass::ipaddr4t
PacketQuery::getIp() const
{
    return ip_;
}

/**************************************************************************************************************/
void
PacketQuery::Stat::add( const Metadata& meta, bool sHit, bool dHit )
{
    if ( sHit )
    {
        nBytesOut_ += meta.nBytes;
    }

    if ( dHit )
    {
        nBytesIn_ += meta.nBytes;
    }

    if ( shouldCountFlows_ )
    {
        BeFlow  flow = { htonl(meta.sip), htonl(meta.dip), htons(meta.vlan), htons(meta.proto), htons(meta.sport), htons(meta.dport) };

        if ( sHit )
        {
            flowsOut_.insert( flow );
            ++nPacketsOut_;
        }

        if ( dHit )
        {
            flowsIn_.insert( flow );
            ++nPacketsIn_;
        }
    }
}

/**************************************************************************************************************/
void
PacketQuery::Stat::countFlows()
{
    shouldCountFlows_ = true;
}

/**************************************************************************************************************/
bool
PacketQuery::Stat::shouldCountFlows() const
{
    return shouldCountFlows_;
}

/**************************************************************************************************************/
SummaryIo
PacketQuery::Stat::getSummary() const
{
    return { flowsIn_.size(), nPacketsIn_, nBytesIn_, flowsOut_.size(), nPacketsOut_, nBytesOut_ };
}
