diff --git a/include/internals.h b/include/internals.h index 7e2ebf2..e1bee6f 100644 --- a/include/internals.h +++ b/include/internals.h @@ -111,6 +111,8 @@ struct enable_if { PDU *pdu_from_flag(Constants::Ethernet::e flag, const uint8_t *buffer, uint32_t size, bool rawpdu_on_no_match = true); +PDU *pdu_from_flag(Constants::IP::e flag, const uint8_t *buffer, + uint32_t size, bool rawpdu_on_no_match = true); PDU *pdu_from_flag(PDU::PDUType type, const uint8_t *buffer, uint32_t size); diff --git a/src/internals.cpp b/src/internals.cpp index d3acfd5..66e04f6 100644 --- a/src/internals.cpp +++ b/src/internals.cpp @@ -34,6 +34,10 @@ #include "radiotap.h" #include "dot11/dot11_base.h" #include "ipv6.h" +#include "tcp.h" +#include "udp.h" +#include "icmp.h" +#include "icmpv6.h" #include "arp.h" #include "eapol.h" #include "rawpdu.h" @@ -100,6 +104,30 @@ Tins::PDU *pdu_from_flag(Constants::Ethernet::e flag, const uint8_t *buffer, }; } +Tins::PDU *pdu_from_flag(Constants::IP::e flag, const uint8_t *buffer, + uint32_t size, bool rawpdu_on_no_match) +{ + switch(flag) { + case Constants::IP::PROTO_IPIP: + return new Tins::TCP(buffer, size); + case Constants::IP::PROTO_TCP: + return new Tins::TCP(buffer, size); + case Constants::IP::PROTO_UDP: + return new Tins::UDP(buffer, size); + case Constants::IP::PROTO_ICMP: + return new Tins::ICMP(buffer, size); + case Constants::IP::PROTO_ICMPV6: + return new Tins::ICMPv6(buffer, size); + case Constants::IP::PROTO_IPV6: + return new Tins::IPv6(buffer, size); + default: + break; + } + if(rawpdu_on_no_match) + return new Tins::RawPDU(buffer, size); + return 0; +} + Tins::PDU *pdu_from_flag(PDU::PDUType type, const uint8_t *buffer, uint32_t size) { switch(type) { diff --git a/src/ip.cpp b/src/ip.cpp index 6acaba7..cd1179e 100644 --- a/src/ip.cpp +++ b/src/ip.cpp @@ -121,30 +121,24 @@ IP::IP(const uint8_t *buffer, uint32_t total_sz) throw malformed_packet(); total_sz -= head_len() * sizeof(uint32_t); if (total_sz) { - switch(_ip.protocol) { - case Constants::IP::PROTO_TCP: - inner_pdu(new Tins::TCP(buffer, total_sz)); - break; - case Constants::IP::PROTO_UDP: - inner_pdu(new Tins::UDP(buffer, total_sz)); - break; - case Constants::IP::PROTO_ICMP: - inner_pdu(new Tins::ICMP(buffer, total_sz)); - break; - case Constants::IP::PROTO_IPV6: - inner_pdu(new Tins::IPv6(buffer, total_sz)); - break; - default: - inner_pdu( - Internals::allocate( - _ip.protocol, - buffer, - total_sz - ) - ); - if(!inner_pdu()) - inner_pdu(new RawPDU(buffer, total_sz)); - break; + inner_pdu( + Internals::pdu_from_flag( + static_cast(_ip.protocol), + buffer, + total_sz, + false + ) + ); + if(!inner_pdu()) { + inner_pdu( + Internals::allocate( + _ip.protocol, + buffer, + total_sz + ) + ); + if(!inner_pdu()) + inner_pdu(new RawPDU(buffer, total_sz)); } } } diff --git a/src/ipv6.cpp b/src/ipv6.cpp index 490fd35..62e79ea 100644 --- a/src/ipv6.cpp +++ b/src/ipv6.cpp @@ -49,6 +49,7 @@ #include "rawpdu.h" #include "exceptions.h" #include "pdu_allocator.h" +#include "internals.h" namespace Tins { @@ -88,30 +89,24 @@ IPv6::IPv6(const uint8_t *buffer, uint32_t total_sz) total_sz -= size; } else { - switch(current_header) { - case Constants::IP::PROTO_TCP: - inner_pdu(new Tins::TCP(buffer, total_sz)); - break; - case Constants::IP::PROTO_UDP: - inner_pdu(new Tins::UDP(buffer, total_sz)); - break; - case Constants::IP::PROTO_ICMP: - inner_pdu(new Tins::ICMP(buffer, total_sz)); - break; - case Constants::IP::PROTO_ICMPV6: - inner_pdu(new Tins::ICMPv6(buffer, total_sz)); - break; - default: - inner_pdu( - Internals::allocate( - current_header, - buffer, - total_sz - ) - ); - if(!inner_pdu()) - inner_pdu(new Tins::RawPDU(buffer, total_sz)); - break; + inner_pdu( + Internals::pdu_from_flag( + static_cast(current_header), + buffer, + total_sz, + false + ) + ); + if(!inner_pdu()) { + inner_pdu( + Internals::allocate( + current_header, + buffer, + total_sz + ) + ); + if(!inner_pdu()) + inner_pdu(new Tins::RawPDU(buffer, total_sz)); } total_sz = 0; }