#include <algorithm>
#include <gtest/gtest.h>
#include "ethernetII.h"
#include "rawpdu.h"
#include "udp.h"
#include "dhcp.h"
#include "dhcpv6.h"
#include "ip.h"

using namespace Tins;

class MatchesResponseTest : public ::testing::Test {
public:
    
};

TEST_F(MatchesResponseTest, TCPSynAck) {
    uint8_t syn_data[] = {
        0, 27, 17, 210, 27, 235, 0, 25, 209, 146, 248, 43, 8, 0, 69, 0, 0, 
        60, 21, 131, 64, 0, 64, 6, 163, 131, 192, 168, 0, 100, 192, 168, 0, 
        1, 219, 85, 31, 144, 11, 209, 99, 140, 0, 0, 0, 0, 160, 2, 57, 8, 
        129, 228, 0, 0, 2, 4, 5, 180, 4, 2, 8, 10, 0, 6, 118, 15, 0, 0, 0, 
        0, 1, 3, 3, 7
    };
    uint8_t resp_data[] = {
        0, 25, 209, 146, 248, 43, 0, 27, 17, 210, 27, 235, 8, 0, 69, 0, 0, 
        40, 0, 0, 64, 0, 64, 6, 185, 26, 192, 168, 0, 1, 192, 168, 0, 100, 
        31, 144, 219, 85, 0, 0, 0, 0, 11, 209, 99, 141, 80, 20, 0, 0, 195, 
        214, 0, 0, 0, 0, 0, 0, 0, 0
    };
    
    EthernetII sent(syn_data, sizeof(syn_data));
    EXPECT_TRUE(sent.matches_response(resp_data, sizeof(resp_data)));
    EXPECT_FALSE(sent.matches_response(syn_data, sizeof(syn_data)));
}

TEST_F(MatchesResponseTest, UDPMatchesICMPResponse) {
    const uint8_t udp_pkt[] = {
        69, 0, 0, 36, 44, 89, 64, 0, 64, 17, 140, 186, 192, 168, 0, 100, 192, 
        168, 0, 1, 136, 167, 39, 15, 0, 16, 129, 215, 74, 69, 74, 69, 74, 69, 
        74, 10
    }, 
    udp_pkt_resp[] = {
        69, 192, 0, 64, 150, 102, 0, 0, 64, 1, 97, 225, 192, 168, 0, 1, 192, 
        168, 0, 100, 3, 3, 126, 212, 0, 0, 0, 0, 69, 0, 0, 36, 44, 89, 64, 0, 
        64, 17, 140, 186, 192, 168, 0, 100, 192, 168, 0, 1, 136, 167, 39, 15, 
        0, 16, 165, 135, 74, 69, 74, 69, 74, 69, 74, 10
    };
    IP ip(udp_pkt, sizeof(udp_pkt));
    EXPECT_TRUE(ip.matches_response(udp_pkt_resp, sizeof(udp_pkt_resp)));
}

TEST_F(MatchesResponseTest, UDPMatchesICMPResponse2) {
    const uint8_t udp_pkt[] = {
        69, 0, 0, 33, 0, 1, 0, 0, 128, 17, 185, 21, 192, 168, 0, 100, 
        192, 168, 0, 1, 0, 0, 13, 5, 0, 13, 62, 59, 98, 111, 105, 110, 
        103
    }, 
    udp_pkt_resp[] = {
        69, 192, 0, 61, 150, 255, 0, 0, 64, 1, 97, 75, 192, 168, 0, 1, 
        192, 168, 0, 100, 3, 3, 126, 209, 0, 0, 0, 0, 69, 0, 0, 33, 0, 
        1, 0, 0, 128, 17, 185, 21, 192, 168, 0, 100, 192, 168, 0, 1, 0, 
        0, 13, 5, 0, 13, 62, 59, 98, 111, 105, 110, 103
    };
    IP ip(udp_pkt, sizeof(udp_pkt));
    EXPECT_TRUE(ip.matches_response(udp_pkt_resp, sizeof(udp_pkt_resp)));
}

TEST_F(MatchesResponseTest, DHCP) {
    uint8_t dhcp_discover[] = {
        255, 255, 255, 255, 255, 255, 0, 1, 1, 0, 0, 1, 8, 0, 69, 0, 1, 
        25, 77, 62, 0, 0, 64, 17, 44, 151, 0, 0, 0, 0, 255, 255, 255, 
        255, 0, 68, 0, 67, 1, 5, 0, 0, 1, 1, 6, 0, 217, 7, 133, 224, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 
        0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 99, 130, 83, 99, 53, 1, 1, 61, 7, 1, 0, 1, 1, 
        0, 0, 1, 255
    };
    
    uint8_t dhcp_offer[] = {
        0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 2, 8, 0, 69, 0, 1, 28, 221, 161, 
        0, 0, 64, 17, 158, 45, 127, 0, 0, 1, 127, 0, 0, 1, 0, 67, 0, 68, 
        1, 8, 0, 0, 2, 1, 6, 0, 217, 7, 133, 224, 0, 0, 0, 0, 0, 0, 0, 0, 
        127, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
        99, 130, 83, 99, 53, 1, 2, 54, 4, 127, 0, 0, 1, 51, 4, 0, 0, 1, 
        44, 255
    };
    
    EthernetII discover(dhcp_discover, sizeof(dhcp_discover));
    UDP* udp = discover.find_pdu<UDP>();
    const RawPDU* raw = discover.find_pdu<RawPDU>();
    ASSERT_TRUE(udp != NULL);
    ASSERT_TRUE(raw != NULL);
    
    udp->inner_pdu(raw->to<DHCP>());
    EXPECT_TRUE(discover.matches_response(dhcp_offer, sizeof(dhcp_offer)));
    EXPECT_FALSE(discover.matches_response(dhcp_discover, sizeof(dhcp_discover)));
}

TEST_F(MatchesResponseTest, ICMP) {
    uint8_t request[] = {
        0, 16, 219, 124, 173, 34, 0, 22, 203, 139, 6, 92, 8, 0, 69, 0, 0, 
        84, 226, 159, 0, 0, 64, 1, 150, 133, 10, 10, 1, 89, 209, 131, 36, 
        158, 8, 0, 3, 45, 54, 12, 0, 0, 146, 91, 68, 72, 241, 31, 12, 0, 
        8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 
        25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 
        41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55
    };
    
    uint8_t reply[] = {
        0, 22, 203, 139, 6, 92, 0, 16, 219, 124, 173, 34, 8, 0, 69, 0, 0, 
        84, 74, 84, 0, 0, 55, 1, 55, 209, 209, 131, 36, 158, 10, 10, 1, 
        89, 0, 0, 11, 45, 54, 12, 0, 0, 146, 91, 68, 72, 241, 31, 12, 0, 
        8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 
        25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 
        41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55
    };
    
    EthernetII pkt(request, sizeof(request));
    EXPECT_TRUE(pkt.matches_response(reply, sizeof(reply)));
    EXPECT_FALSE(pkt.matches_response(request, sizeof(request)));
}

TEST_F(MatchesResponseTest, ARP) {
    uint8_t request[] = {
        255, 255, 255, 255, 255, 255, 0, 1, 1, 0, 0, 1, 8, 6, 0, 1, 8, 0, 
        6, 4, 0, 1, 0, 1, 1, 0, 0, 1, 127, 0, 0, 1, 255, 255, 255, 255, 
        255, 255, 127, 0, 0, 1
    };
    
    uint8_t reply[] = {
        0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 2, 8, 6, 0, 1, 8, 0, 6, 4, 0, 2, 
        0, 1, 1, 0, 0, 2, 127, 0, 0, 1, 0, 1, 1, 0, 0, 1, 127, 0, 0, 1
    };
    
    EthernetII pkt(request, sizeof(request));
    EXPECT_TRUE(pkt.matches_response(reply, sizeof(reply)));
    EXPECT_FALSE(pkt.matches_response(request, sizeof(request)));
}

TEST_F(MatchesResponseTest, ICMPv6) {
    uint8_t request[] = {
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 134, 221, 96, 0, 0, 0, 0, 64, 
        58, 64, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 128, 0, 91, 104, 25, 156, 
        0, 1, 226, 206, 89, 81, 0, 0, 0, 0, 14, 139, 1, 0, 0, 0, 0, 0, 16, 
        17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 
        33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 
        49, 50, 51, 52, 53, 54, 55
    };
    
    uint8_t reply[] = {
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 134, 221, 96, 0, 0, 0, 0, 64, 
        58, 64, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 129, 0, 90, 104, 25, 156, 
        0, 1, 226, 206, 89, 81, 0, 0, 0, 0, 14, 139, 1, 0, 0, 0, 0, 0, 
        16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 
        32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 
        48, 49, 50, 51, 52, 53, 54, 55
    };
    
    uint8_t not_a_reply[] = {
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 134, 221, 96, 0, 0, 0, 0, 26, 
        60, 64, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 17, 0, 11, 1, 9, 1, 1, 0, 
        166, 17, 0, 42, 0, 18, 104, 137, 97, 98, 99, 100, 101, 102, 103, 
        104, 105, 106
    };
    
    EthernetII pkt(request, sizeof(request));
    EXPECT_TRUE(pkt.matches_response(reply, sizeof(reply)));
    EXPECT_FALSE(pkt.matches_response(request, sizeof(request)));
    EXPECT_FALSE(pkt.matches_response(not_a_reply, sizeof(not_a_reply)));
}

TEST_F(MatchesResponseTest, DHCPv6) {
    uint8_t request[] = {
        51, 51, 0, 1, 0, 2, 0, 2, 179, 193, 64, 207, 134, 221, 96, 0, 0, 
        0, 0, 52, 17, 128, 254, 128, 0, 0, 0, 0, 0, 0, 2, 2, 179, 255, 
        254, 193, 64, 207, 255, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 
        0, 2, 2, 34, 2, 35, 0, 52, 243, 80, 1, 0, 83, 101, 0, 1, 0, 14, 
        0, 1, 0, 6, 17, 195, 100, 52, 0, 22, 118, 189, 211, 18, 0, 3, 0, 
        12, 0, 0, 0, 1, 255, 255, 255, 255, 255, 255, 255, 255, 0, 8, 0, 
        2, 0, 100
    };
    
    uint8_t reply[] = {
        0, 2, 179, 193, 64, 207, 0, 21, 242, 4, 141, 219, 134, 221, 96, 0, 
        0, 0, 0, 209, 17, 128, 254, 128, 0, 0, 0, 0, 0, 0, 2, 21, 242, 255, 
        254, 4, 141, 219, 254, 128, 0, 0, 0, 0, 0, 0, 2, 2, 179, 255, 254, 
        193, 64, 207, 2, 35, 2, 34, 0, 209, 127, 59, 2, 0, 83, 101, 0, 1, 
        0, 14, 0, 1, 0, 6, 17, 195, 100, 52, 0, 22, 118, 189, 211, 18, 0, 
        3, 0, 132, 0, 0, 0, 1, 0, 0, 14, 16, 0, 0, 21, 24, 0, 5, 0, 24, 
        102, 0, 0, 0, 0, 0, 0, 0, 213, 240, 138, 111, 157, 131, 93, 191, 
        0, 1, 81, 128, 0, 2, 163, 0, 0, 13, 0, 88, 0, 0, 49, 32, 97, 100, 
        100, 114, 101, 115, 115, 32, 103, 114, 97, 110, 116, 101, 100, 46, 
        32, 89, 111, 117, 32, 109, 97, 121, 32, 105, 110, 99, 108, 117, 
        100, 101, 32, 73, 65, 65, 68, 68, 82, 32, 105, 110, 32, 73, 65, 
        32, 111, 112, 116, 105, 111, 110, 44, 32, 105, 102, 32, 121, 111, 
        117, 32, 119, 97, 110, 116, 32, 116, 111, 32, 112, 114, 111, 118, 
        105, 100, 101, 32, 97, 32, 104, 105, 110, 116, 46, 0, 2, 0, 14, 0, 
        1, 0, 1, 17, 195, 197, 84, 0, 224, 129, 73, 16, 201, 0, 7, 0, 1, 
        0, 0, 12, 0, 16, 0, 119, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6
    };
    
    EthernetII pkt(request, sizeof(request));
    UDP* udp = pkt.find_pdu<UDP>();
    const RawPDU* raw = pkt.find_pdu<RawPDU>();
    ASSERT_TRUE(udp != NULL);
    ASSERT_TRUE(raw != NULL);
    
    udp->inner_pdu(raw->to<DHCPv6>());
    EXPECT_TRUE(pkt.matches_response(reply, sizeof(reply)));
    EXPECT_FALSE(pkt.matches_response(request, sizeof(request)));
}