// Boost pinger (c) 2011 by Guilherme Maciel Ferreira / Intra2net AG
// Based upon work copyright (c) 2003-2010 Christopher M. Kohlhoff (ping.cpp)
//
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)
#include "icmp/icmppinger.h"

#include <errno.h>

#include <ostream>

#include <boost/bind.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <boost/date_time/posix_time/posix_time_types.hpp>
#include <boost/uuid/uuid.hpp>
#include <boost/uuid/uuid_generators.hpp>
#include <boost/foreach.hpp>
#include <boost/system/system_error.hpp>

#include <logfunc.hpp>

#include "boost_assert_handler.h"
#include "icmp/icmppacketfactory.h"
#include "host/networkinterface.hpp"

using namespace std;
using boost::asio::const_buffers_1;
using boost::asio::io_service;
using boost::asio::ip::address;
using boost::asio::ip::icmp;
using boost::function;
using boost::posix_time::microsec_clock;
using boost::posix_time::seconds;
using boost::shared_ptr;
using I2n::Logger::GlobalLogger;

using boost::asio::ip::icmp;

//-----------------------------------------------------------------------------
// IcmpPinger
//-----------------------------------------------------------------------------

/**
 * @brief factory function for IcmpPingers, ensures that set_myself is set
 *
 * @returns a shared pointer to a Pinger
 */
PingerItem IcmpPinger::create(
        const IoServiceItem io_serv,
        const icmp::socket::protocol_type &protocol,
        const string &source_network_interface,
        const int echo_reply_timeout_in_sec )
{
    // get distributor
    IcmpPacketDistributorItem distributor = IcmpPacketDistributor::get_distributor(
            protocol, source_network_interface, io_serv);

    // create pinger
    IcmpPinger *ptr = new IcmpPinger(io_serv, protocol, echo_reply_timeout_in_sec, distributor);
    IcmpPingerItem shared_ptr_(ptr);

    // keep weak pointer to self
    //shared_ptr_->set_myself( weak_ptr ); //Error: Pinger::set_myself is protected
    ptr->set_myself( shared_ptr_ );

    // register in distributor
    distributor->register_pinger(shared_ptr_);

    // done, return shared ptr
    return shared_ptr_;
}

/**
 * @brief Parameterized constructor.
 *
 * @param io_serv The one @c io_service object that controls async processing
 * @param protocol The network layer protocol to use.
 * @param source_network_interface The network interface name from where to
 * send the packets.
 * @param echo_reply_timeout_in_sec The amount of time to wait for a reply.
 */
IcmpPinger::IcmpPinger(
        const IoServiceItem io_serv,
        const icmp::socket::protocol_type &protocol,
        const int echo_reply_timeout_in_sec,
        const IcmpPacketDistributorItem distributor
) :
    PacketDistributor( distributor ),
    DestinationEndpoint(),
    Protocol( protocol ),
    IcmpPacketReceiveTimer( *io_serv ),
    Identifier( 0 ),
    SequenceNumber( 0 ),
    TimeSent( microsec_clock::universal_time() ),
    ReplyReceived( false ),
    EchoReplyTimeoutInSec( echo_reply_timeout_in_sec ),
    PingerStatus( PingStatus_NotSent ),
    PingDoneCallback(),
    LogPrefix("IcmpPinger")
{
    // Create "unique" identifier
    boost::uuids::random_generator random_gen;
    boost::uuids::uuid random_tag = random_gen();

    BOOST_ASSERT( sizeof(Identifier) <= random_tag.size() );
    memcpy( &Identifier, random_tag.data, sizeof(Identifier) );

    LogPrefix = "IPing(no IP yet): ";
}

/**
 * @brief Destructor.
 */
IcmpPinger::~IcmpPinger()
{
}

/**
 * @brief Ping a destination address from an available local source.
 *
 * @param destination_ip The address of the host to ping.
 * @param destination_port The port at the destination host to ping.
 * @param done_handler Done handler will be called on successful ping or timeout.
 *
 * @return void.
 */
void IcmpPinger::ping(
        const address &destination_ip,
        const uint16_t /*destination_port*/, // the ICMP protocol does not use ports
        function<void(PingStatus,long)> ping_done_callback
)
{
    PingDoneCallback = ping_done_callback;

    // Prepare ping
    set_ping_status( PingStatus_NotSent );

    set_destination_endpoint( destination_ip );

    start_send();
}

void IcmpPinger::stop_pinging()
{
    GlobalLogger.debug() << LogPrefix << "stop_pinging" << endl;

    GlobalLogger.debug() << LogPrefix << "cancel timer" << endl;
    IcmpPacketReceiveTimer.cancel();

    GlobalLogger.debug() << LogPrefix << "unregister" << endl;

    IcmpPingerItem icmp_item = boost::static_pointer_cast<IcmpPinger>( get_myself().lock() );
    if ( icmp_item )
    {
        PacketDistributor->unregister_pinger( icmp_item );
    } else
    {
        GlobalLogger.warning() << LogPrefix
            << "weak pointer to pinger broken is empty. Huh?" << endl;
    }
}


void IcmpPinger::set_destination_endpoint( const address &destination_ip )
{
    uint16_t port = 0;
    DestinationEndpoint = icmp::endpoint( destination_ip, port );

    // update LogPrefix
    std::stringstream temp;
    temp << "IPing(" << DestinationEndpoint.address().to_string() << "): ";
    LogPrefix = temp.str();
}

bool IcmpPinger::start_send()
{
    ++SequenceNumber;

    IcmpPacketItem icmp_packet_echo_request = IcmpPacketFactory::create_icmp_packet_echo_request(
            Protocol, Identifier, SequenceNumber );

    BOOST_ASSERT( PingerStatus == PingStatus_NotSent );
    return send_echo_request( icmp_packet_echo_request );
}

bool IcmpPinger::send_echo_request( const IcmpPacketItem icmp_packet )
{
    boost::asio::streambuf request_buffer;
    ostream os( &request_buffer );
    if ( !icmp_packet->write( os ) )
    {
        GlobalLogger.error() << LogPrefix << "fail writing ping data." << endl;
    }

    TimeSent = microsec_clock::universal_time();

    string dest_address_string = DestinationEndpoint.address().to_string();
    BOOST_ASSERT( !dest_address_string.empty() );

    // Send the request
    size_t bytes_sent = 0;
    try
    {
        GlobalLogger.info() << LogPrefix << "sending ping" << endl;
        const_buffers_1 data = request_buffer.data();

        // Block until send the data
        bytes_sent = PacketDistributor->get_socket()->send_to( data, DestinationEndpoint );
        if ( bytes_sent != buffer_size( data ) )
        {
            GlobalLogger.error() << LogPrefix << "fail sending ping data. Only"
                                 << bytes_sent << " of " << buffer_size(data)
                                 << " bytes were sent!" << endl;
        }

        ReplyReceived = false;
        schedule_timeout_echo_reply();
    }
    catch ( const boost::system::system_error &boost_err )
    {
        boost::system::error_code err_code = boost_err.code();
        GlobalLogger.error() << LogPrefix << "fail sending ping data: "
                 << boost_err.what() << " (code " << err_code.value()
                 << ", category " << err_code.category().name() << ")" << endl;

        // do not wait for timeout but fail at once
        set_ping_status(PingStatus_SendFailed);
        ReplyReceived = true;   // flag for handler to leave ping status as is
        //handle_timeout( err_code );
        handle_timeout( boost::system::error_code() );
    }
    catch ( const exception &ex )
    {
        GlobalLogger.error() << LogPrefix << "fail sending ping data: "
                             << ex.what() << endl;

        // do not wait for timeout but fail at once
        set_ping_status(PingStatus_SendFailed);
        ReplyReceived = true;   // flag for handler to leave ping status as is
        handle_timeout( boost::system::error_code() );
    }
    catch ( ... )
    {
        GlobalLogger.error() << LogPrefix << "fail sending ping data: "
                             << "Unknown exception" << endl;

        // do not wait for timeout but fail at once
        set_ping_status(PingStatus_SendFailed);
        ReplyReceived = true;   // flag for handler to leave ping status as is
        handle_timeout( boost::system::error_code() );
    }

    return (bytes_sent > 0);
}

void IcmpPinger::schedule_timeout_echo_reply()
{
    // Wait up to N seconds for a reply.
    (void) IcmpPacketReceiveTimer.expires_at(
            TimeSent + seconds( EchoReplyTimeoutInSec )
    );
    IcmpPacketReceiveTimer.async_wait(
            boost::bind( &IcmpPinger::handle_timeout, this, boost::asio::placeholders::error )
    );
}

/**
 * @brief Gets called when the ping is finished: Either on timeout or on ping reply
 *
 * @return void (but calls PingDoneCallback)
 **/
void IcmpPinger::handle_timeout(const boost::system::error_code& error)
{
    if (error)
    {
        if ( error ==  boost::asio::error::operation_aborted )
        {
            if (! ReplyReceived)
            {
                GlobalLogger.notice() << LogPrefix
                    << "Timer waiting for ICMP echo reply was cancelled!"
                    << endl;
                set_ping_status( PingStatus_FailureAsyncCancel );
            }
            // otherwise probably called by IcmpPacketReceiveTimer.cancel in
            // handle_receive_icmp_packet!
        }
        else
        {
            GlobalLogger.notice() << LogPrefix << "Error " << error
                << " waiting for ICMP echo reply!" << endl;
            set_ping_status( PingStatus_FailureAsyncError );
        }
        // could check here for more details if error is forwarded from
        // send_echo_request

        // Still continue with rest of function, so PingStatus is updated and Callback executed
        //   when timer was cancelled
    }
    else if ( !ReplyReceived )
    {    // Check ReplyReceived since the timer handler is also called by Timer.cancel();
        GlobalLogger.info() << LogPrefix << "Request timed out" << endl;

        set_ping_status( PingStatus_FailureTimeout );
    }
    // otherwise assume that ping status was set already

    // Call ping-done handler
    PingDoneCallback( PingerStatus, static_cast<long>(
                     (microsec_clock::universal_time()
                                            - TimeSent).total_microseconds()) );
}


/**
 * @brief Receive ICMP packets
 * @param bytes_transferred Number of bytes transferred.
 * @return true if packet matches a request from this pinger, false otherwise
 **/
bool IcmpPinger::handle_receive_icmp_packet(const IcmpPacketItem icmp_packet,
                                            const size_t bytes_transferred )
{
    bool does_match = false;

    if ( ReplyReceived )
    {
        // continue, might be an old packet
        // or return false right away, do not want packet anyway...
        return does_match;
    }
    else if ( DestinationEndpoint.address() == address() )
    {   // we have no IP set yet
        return does_match;
    }

    // We can receive all ICMP packets received by the host, so we need to
    // filter out only the echo replies that match our identifier,
    // expected sequence number, and destination host address (receive just
    // the ICMP packets from the host we had ping).

    try
    {
        if ( icmp_packet->match_echo_reply(
                                Identifier, SequenceNumber,
                                DestinationEndpoint.address() ) )
        {
            GlobalLogger.info() << LogPrefix << "Received reply" << endl;

            ReplyReceived = true;
            does_match = true;

            icmp_packet->print( bytes_transferred, TimeSent );

            set_ping_status( PingStatus_SuccessReply );

            IcmpPacketReceiveTimer.cancel();                            //lint !e534
        }
        else if ( icmp_packet->match_destination_unreachable(
                                     Identifier, SequenceNumber,
                                     DestinationEndpoint.address() ) )
        {
            GlobalLogger.info() << LogPrefix
                                << "Received destination unreachable" << endl;

            ReplyReceived = true;
            does_match = true;

            icmp_packet->print( bytes_transferred, TimeSent );

            set_ping_status( PingStatus_FailureDestinationUnreachable );

            IcmpPacketReceiveTimer.cancel();                            //lint !e534
        }
        else if ( icmp_packet->match_time_exceeded(
                                     Identifier, SequenceNumber,
                                     DestinationEndpoint.address() ) )
        {
            GlobalLogger.info() << LogPrefix
                                << "Received time exceeded" << endl;

            ReplyReceived = true;
            does_match = true;

            icmp_packet->print( bytes_transferred, TimeSent );

            set_ping_status( PingStatus_FailureDestinationUnreachable );

            IcmpPacketReceiveTimer.cancel();                            //lint !e534
        }
        else
        {
            GlobalLogger.debug() << LogPrefix
               << "Received packet that does not match or has wrong seq.nr"
               << endl;
        }
    }
    catch ( std::exception &exc)
    {
        GlobalLogger.warning() << LogPrefix
            << "Caught exception in packet interpretation: " << exc.what()
            << std::endl;
        if ( IcmpPacketFactory::PacketDumpMode == DUMP_ALWAYS ||
             IcmpPacketFactory::PacketDumpMode == DUMP_IF_ERROR )
            IcmpPacketFactory::dump_packet(*icmp_packet);
        does_match = true;   // avoid the same procedure in all other pingers
    }
    catch ( ... )
    {
        GlobalLogger.warning() << LogPrefix
            << "Caught unspecified exception in packet interpretation!"
            << std::endl;
        if ( IcmpPacketFactory::PacketDumpMode == DUMP_ALWAYS ||
             IcmpPacketFactory::PacketDumpMode == DUMP_IF_ERROR )
            IcmpPacketFactory::dump_packet(*icmp_packet);
        does_match = true;   // avoid the same procedure in all other pingers
    }

    return does_match;
}

void IcmpPinger::set_ping_status( PingStatus ping_status )
{
    PingerStatus = ping_status;
}

//------------------------------------------------------------------------
// IcmpPacketDistributor
//------------------------------------------------------------------------

static const std::size_t SOCKET_BUFFER_SIZE = 65536;   // 64kB

typedef std::set<IcmpPingerItem>::iterator PingerListIterator;


bool IcmpPacketDistributor::InstanceIdentifierComparator::operator() (
                const IcmpPacketDistributor::DistributorInstanceIdentifier &a,
                const IcmpPacketDistributor::DistributorInstanceIdentifier &b )
                                                                          const
{
    if ( a.first == boost::asio::ip::icmp::v4() )
    {
        if ( b.first == boost::asio::ip::icmp::v4() )
            return a.second < b.second;   // v4 == v4
        else
            BOOST_ASSERT( b.first == boost::asio::ip::icmp::v6() );
            return true;    // a(v4) < b(b6)
    }
    else
    {
        BOOST_ASSERT( a.first == boost::asio::ip::icmp::v6() );

        if ( b.first == boost::asio::ip::icmp::v4() )
            return false;   // a(v6) > b(v4)
        else
            BOOST_ASSERT( b.first == boost::asio::ip::icmp::v6() );
            return a.second < b.second;    // v6 == v6
    }
}

//-----------------------------------------------------------------------------
// Definition of IcmpPacketDistributor
//-----------------------------------------------------------------------------

IcmpPacketDistributor::map_type IcmpPacketDistributor::Instances; // initialize


IcmpPacketDistributorItem IcmpPacketDistributor::get_distributor(
        const icmp::socket::protocol_type &protocol,
        const std::string &network_interface,
        const IoServiceItem io_serv )
{
    IcmpPacketDistributor::DistributorInstanceIdentifier identifier(
                                                  protocol, network_interface);

    // check if there is an instance for this protocol and interface
    if ( Instances.count(identifier) == 0 )
    {   // need to create an instance for this protocol and network interface
        std::string protocol_str;
        if (protocol == icmp::v4())
            protocol_str = "ICMPv4";
        else if (protocol == icmp::v6())
            protocol_str = "ICMPv6";
        else
            protocol_str = "unknown protocol!";

        GlobalLogger.info() << "Creating IcmpPacketDistributor for interface "
                            << network_interface << " and protocol "
                            << protocol_str << std::endl;
        IcmpPacketDistributorItem new_instance( new IcmpPacketDistributor(
                    protocol, network_interface, io_serv ) );
        Instances[identifier] = new_instance;
    }

    BOOST_ASSERT( Instances.count(identifier) == 1 );

    // return the one instance for this protocol and interface
    return Instances[identifier];
}


IcmpPacketDistributorItem IcmpPacketDistributor::get_distributor(
        const icmp::socket::protocol_type &protocol,
        const std::string &network_interface )
{
    IcmpPacketDistributor::DistributorInstanceIdentifier identifier(
                                                  protocol, network_interface);

    BOOST_ASSERT( Instances.count(identifier) == 1 );

    // return the one instance for this protocol and interface
    return Instances[identifier];
}


IcmpPacketDistributor::IcmpPacketDistributor(
            const icmp::socket::protocol_type &protocol,
            const std::string &network_interface,
            const IoServiceItem io_serv ):
    Protocol( protocol ),
    Socket( new icmp::socket(*io_serv, protocol) ),
    ReplyBuffer(),
    PingerList()
{
    // set TTL for testing
    //const boost::asio::ip::unicast::hops option( 3 );
    //Socket->set_option(option);

    NetworkInterface<icmp::socket, boost::asio::ip::icmp>
                  NetInterface( network_interface, *Socket );

    if ( !NetInterface.bind() )
    {
        GlobalLogger.error()
           << "Trouble creating IcmpPacketDistributor for interface "
           << network_interface// << " and protocol " << protocol
           << ": could not bind the socket with the local interface. "
           << ::strerror( errno )  << std::endl;
    }

    register_receive_handler();
}


void IcmpPacketDistributor::register_receive_handler()
{
    // wait for reply, prepare buffer to receive up to SOCKET_BUFFER_SIZE bytes
    Socket->async_receive(
            ReplyBuffer.prepare( SOCKET_BUFFER_SIZE ),
            boost::bind( &IcmpPacketDistributor::handle_receive, this,
                         boost::asio::placeholders::error,
                         boost::asio::placeholders::bytes_transferred )
    );
}

void IcmpPacketDistributor::handle_receive(
                                        const boost::system::error_code &error,
                                        const size_t &bytes_transferred )
{
    if ( error )
    {
        GlobalLogger.warning()
           << ": Received error " << error
           << " in ICMP packet distributor; end handler and schedule another.";
        register_receive_handler();
        return;
    }

    // The actual number of bytes received is committed to the buffer so that we
    // can extract it using a std::istream object.
    ReplyBuffer.commit( bytes_transferred );

    GlobalLogger.info() << "received packet in distributor" << std::endl;

    std::istream is( &ReplyBuffer );
    if ( !is )
    {
        GlobalLogger.error() << "Can't handle ReplyBuffer" << std::endl;
        return;
    }

    // Decode the reply packet.
    IcmpPacketItem icmp_packet = IcmpPacketFactory::create_icmp_packet(
                                                             Protocol, is );
    if ( !icmp_packet )
    {
        GlobalLogger.warning() << "Ignoring broken ICMP packet"
                               << std::endl;
    }
    else
    {
        GlobalLogger.debug() << "Succesfully parsed ICMP packet"
                             << std::endl;

        // check which pinger wants this packet
        bool packet_matches = false;
        BOOST_FOREACH( const IcmpPingerItem &pinger, PingerList )
        {
            packet_matches = pinger->handle_receive_icmp_packet(
                                            icmp_packet, bytes_transferred);
            if (packet_matches)
                break;
        }
        if (!packet_matches)
            GlobalLogger.info() << "Packet did not match any pinger"
                                << std::endl;
    }

    // re-register receive handler
    register_receive_handler();
}

bool IcmpPacketDistributor::register_pinger( const IcmpPingerItem &new_pinger )
{
    std::pair<PingerListIterator, bool> result = PingerList.insert(new_pinger);
    bool was_new = result.second;
    if (was_new)
        GlobalLogger.info() << "Register new pinger with IcmpPacketDistributor"
                            << std::endl;
    else
        GlobalLogger.warning()
            << "Pinger to register was already known in IcmpPacketDistributor"
            << std::endl;
    return was_new;
}


bool IcmpPacketDistributor::unregister_pinger( const IcmpPingerItem &old_pinger )
{
    int n_erased = PingerList.erase(old_pinger);
    bool was_erased = n_erased > 0;
    if (was_erased)
        GlobalLogger.info() << "Removed pinger from IcmpPacketDistributor"
                            << std::endl;
    else
        GlobalLogger.warning()
            << "Could not find pinger to remove from IcmpPacketDistributor"
            << std::endl;
    return was_erased;
}

/**
 * @brief for all instances: close sockets, unregister all pingers
 */
void IcmpPacketDistributor::clean_up_all()
{
    BOOST_FOREACH( IcmpPacketDistributor::map_type::value_type &instance,
                                                                    Instances )
    {
        instance.second->clean_up();
    }

    Instances.clear();
}

void IcmpPacketDistributor::clean_up()
{
    if (PingerList.size() == 0)
        GlobalLogger.info() << "All IcmpPingers have de-registered"
                            << std::endl;
    else
        GlobalLogger.warning() << "There were still " << PingerList.size()
            << " pingers registered in IcmpPacketDistributor!" << std::endl;
    PingerList.clear();

    boost::system::error_code error;
    //Socket->shutdown(icmp::socket::shutdown_both, error);  //both=send&receive
    //if ( error )
    //    GlobalLogger.warning() << "Received error " << error
    //                           << " when shutting down ICMP socket";
    // always gave an error system:9 (probably EBADF: Bad file descriptor)

    Socket->close(error);
    if ( error )
        GlobalLogger.warning() << "Received error " << error
                               << " when closing ICMP socket";
}

IcmpPacketDistributor::~IcmpPacketDistributor()
{
    GlobalLogger.info() << "Destroying IcmpPacketDistributor" << std::endl;
}

SocketItem IcmpPacketDistributor::get_socket() const
{
    return Socket;
}