/* namecollection.cc
****************************************************************************************************************
****************************************************************************************************************

    Copyright (C) 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2020, 2021, 2026 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 <cassert>
#include <cctype>
#include <vector>

#include "namecollection.h"
#include "pcapwriting.h"

using namespace iclass;
using namespace trafficmeter;

using std::map;
using std::string;
using std::transform;
using std::vector;

constexpr const char*   DEBUG_FILE_     = "debug.pcap";
constexpr const char*   VALID_CHAR_     = "-.abcdefghijklmnopqrstuvwxyz0123456789";
constexpr const char*   ALNUM_CHAR_     = VALID_CHAR_ + 2;

static const char*      EXCLUSION_NETWORKS_[] =
{
    "0/8",
    "10/8",
    "100.64/10",
    "127/8",
    "169.254/16",
    "172.16/12",
    "192.168/16",
    "224/4",
    "255.255.255.255"
};

/**************************************************************************************************************/
static bool
trimName_( string& name )
{
    transform( name.cbegin(), name.cend(), name.begin(), tolower );
    name.erase( name.find_last_of(ALNUM_CHAR_) + 1 ).erase( 0, name.find_first_of(ALNUM_CHAR_) );
    return !name.empty() && name.find_first_not_of(VALID_CHAR_) == string::npos;
}

/**************************************************************************************************************/
static void
debugOut_( Packet packet )
{
    try
    {
        PcapWriting( DEBUG_FILE_ ).writePacket( packet.timestamp(), packet.data2(), packet.nBytes2() );
    }
    catch ( ... )
    {
    }
}

/**************************************************************************************************************/
static vector<NameCollection::Key>
extract_( Packet packet, const Metadata& meta )
{
    vector<NameCollection::Key>     names;
    const uint8_t*                  dns = packet.data2() + meta.offsetLayer4 + 8;
    const uint8_t*                  end = packet.data2() + packet.nBytes2();

    if
    (
        !
        (
            (dns[ 2 ] & 0xf8) == 0x80 &&        // QR is response and OPCODE is query
            (dns[ 3 ] & 0x0f) == 0 &&           // RCODE is no error
            end - dns >= 35                     // enough packet size
        )
    )
    {
        return std::move( names );
    }

    int     nQuestions = ntohs( *reinterpret_cast<const uint16_t*>(dns + 4) );
    int     nAnswers = ntohs( *reinterpret_cast<const uint16_t*>(dns + 6) );
    int     nAuthorities = ntohs( *reinterpret_cast<const uint16_t*>(dns + 8) );
    int     nAdditionals = ntohs( *reinterpret_cast<const uint16_t*>(dns + 10) );

    if ( nAnswers == 0 && nAdditionals == 0 )
    {
        return std::move( names );
    }

    if ( nQuestions != 1 )  // FORMERR
    {
        return std::move( names );
    }

    const uint8_t*  question = dns + 12;
    const uint8_t*  p = reinterpret_cast<const uint8_t*>( memchr(question, 0, end-question) );

    if ( p == nullptr )
    {
        return std::move( names );
    }

    uint16_t    qtype = ntohs( *reinterpret_cast<const uint16_t*>(p + 1) );
    uint16_t    qclass = ntohs( *reinterpret_cast<const uint16_t*>(p + 3) );

    if ( qclass != 1 )
    {
        return std::move( names );
    }

    p += 5;

    if ( qtype == 1 )   // IN A answer
    {
        string  qname( reinterpret_cast<const char*>(question + 1) );

        for ( unsigned i = *question; question[i + 1] > 0; i += question[i + 1] + 1 )
        {
            if ( i >= qname.size() )
            {
                debugOut_( packet );
                return std::move( names );
            }

            qname[ i ] = '.';
        }

        if ( !trimName_(qname) )    // QNAME is invalid
        {
            return std::move( names );
        }

        for ( int i = 0; i < nAnswers; ++i )
        {
            if ( end - p < 16 )
            {
                return std::move( names );
            }

            while ( true )
            {
                if ( *p == 0 )
                {
                    p++;
                    break;
                }

                if ( (*p & 0xc0) == 0xc0 )
                {
                    p += 2;
                    break;
                }

                p += *p + 1;

                if ( p >= end )
                {
                    return std::move( names );
                }
            }

            if ( end - p < 14 )
            {
                return std::move( names );
            }

            uint16_t    rtype = ntohs( *reinterpret_cast<const uint16_t*>(p) );
            uint16_t    rclass = ntohs( *reinterpret_cast<const uint16_t*>(p + 2) );
            uint16_t    rdatalen = ntohs( *reinterpret_cast<const uint16_t*>(p + 8) );

            p += 10;

            if
            (
                rtype == 1 &&   // TYPE is A
                rclass == 1 &&  // CLASS is IN
                rdatalen == 4   // RDATA is IPv4
            )
            {
                names.emplace_back( ntohl(*reinterpret_cast<const uint32_t*>(p)), qname );
            }

            p += rdatalen;
        }
    }
    else if ( qtype == 15 ) // IN MX answer
    {
        for ( int i = 0; i < nAnswers; ++i )
        {
            if ( end - p < 12 )
            {
                return std::move( names );
            }

            while ( true )
            {
                if ( *p == 0 )
                {
                    p++;
                    break;
                }

                if ( (*p & 0xc0) == 0xc0 )
                {
                    p += 2;
                    break;
                }

                p += *p + 1;

                if ( p >= end )
                {
                    return std::move( names );
                }
            }

            if ( end - p < 10 )
            {
                return std::move( names );
            }

            p += 10 + ntohs( *reinterpret_cast<const uint16_t*>(p + 8) );
        }

        for ( int i = 0; i < nAuthorities; ++i )
        {
            if ( end - p < 12 )
            {
                return std::move( names );
            }

            while ( true )
            {
                if ( *p == 0 )
                {
                    p++;
                    break;
                }

                if ( (*p & 0xc0) == 0xc0 )
                {
                    p += 2;
                    break;
                }

                p += *p + 1;

                if ( p >= end )
                {
                    return std::move( names );
                }
            }

            if ( end - p < 10 )
            {
                return std::move( names );
            }

            p += 10 + ntohs( *reinterpret_cast<const uint16_t*>(p + 8) );
        }

        for ( int i = 0; i < nAdditionals; ++i )
        {
            if ( end - p < 16 )
            {
                return std::move( names );
            }

            string  rname;

            while ( true )
            {
                if ( *p == 0 )
                {
                    p++;
                    break;
                }

                if ( (*p & 0xc0) == 0xc0 )
                {
                    const uint8_t*  rp = dns + ( ntohs(*reinterpret_cast<const uint16_t*>(p)) & 0x3fff );

                    if ( end - rp < 1 )
                    {
                        return std::move( names );
                    }

                    while ( *rp )
                    {
                        if ( (*rp & 0xc0) == 0xc0 )
                        {
                            rp = dns + ( ntohs(*reinterpret_cast<const uint16_t*>(rp)) & 0x3fff );
                            continue;
                        }

                        if ( end - rp < *rp + 1 )
                        {
                            return std::move( names );
                        }

                        rname += '.';
                        rname.append( reinterpret_cast<const char*>(rp + 1), *rp );
                        rp += *rp + 1;
                    }

                    p += 2;
                    break;
                }

                rname += '.';
                rname.append( reinterpret_cast<const char*>(p + 1), *p );
                p += *p + 1;

                if ( p >= end )
                {
                    return std::move( names );
                }
            }

            if ( end - p < 14 )
            {
                return std::move( names );
            }

            uint16_t    rtype = ntohs( *reinterpret_cast<const uint16_t*>(p) );
            uint16_t    rclass = ntohs( *reinterpret_cast<const uint16_t*>(p + 2) );
            uint16_t    rdatalen = ntohs( *reinterpret_cast<const uint16_t*>(p + 8) );

            p += 10;

            if
            (
                rtype == 1 &&           // TYPE is A
                rclass == 1 &&          // CLASS is IN
                rdatalen == 4 &&        // RDATA is IPv4
                trimName_( rname )      // RNAME is valid
            )
            {
                names.emplace_back( ntohl(*reinterpret_cast<const uint32_t*>(p)), rname );
            }

            p += rdatalen;
        }
    }

    return std::move( names );
}

/**************************************************************************************************************/
NameCollection::NameCollection()
{
    ipset4t  networks;

    for ( const char* n : EXCLUSION_NETWORKS_ )
    {
        networks.insert( ipnet4t(n) );
    }

    exclusions_ = new iplookup4t( networks );
    map_c_ = new map<Key,int>;
}

/**************************************************************************************************************/
NameCollection::~NameCollection()
{
    this->lock();
    assert( map_c_->empty() );
    delete map_c_;
    delete exclusions_;
}

/**************************************************************************************************************/
void
NameCollection::push( const Metadata& meta, Packet packet )
{
    for ( const auto& key : extract_(packet, meta) )
    {
        if ( !exclusions_->contains(key.ip) )
        {
            this->lock();

            auto&&  [ row, wasInserted ] = map_c_->try_emplace( key, 1 );

            if ( !wasInserted )
            {
                row->second++;
            }

            this->unlock();
        }
    }
}

/**************************************************************************************************************/
map<NameCollection::Key,int>*
NameCollection::pull()
{
    map<Key,int>*  pullMap;

    this->lock();
    pullMap = map_c_;
    map_c_ = new map<Key,int>;
    this->unlock();
    return pullMap;
}

/**************************************************************************************************************/
long
NameCollection::nRows() const
{
    this->lock();

    long  nRows = map_c_->size();

    this->unlock();
    return nRows;
}
