#!/usr/bin/python3

#from __future__ import unicode_literals
import ipaddress
from scapy.all import *
import array
import struct

# stolen from scapy (little endian system)
def checksum_scapy(pkt):
    if len(pkt) % 2 == 1:
        pkt += b"\0"
    s = sum(array.array("H", pkt))
    s = (s >> 16) + (s & 0xffff)
    s += s >> 16
    s = ~s
    return (((s>>8)&0xff)|s<<8) & 0xffff

# a) Compare for TCP, UDP -> IPv6 does not have checksum!
# 1. convert to array of "bytes"
# 2. import into an array
# 3. sum everything up
# b) Generate checksum from v4 offset and IPv6 for IPv4

# a)
# UDP/TCP CONTENT stays the same
# Only the diff between v4 and v6 "counts"
#

def sum_for_udp(packet):
    sums = ""
    sums += struct.pack("H", packet[UDP].sport)
    sums += struct.pack("H", packet[UDP].dport)
    sums += struct.pack("H", packet[UDP].len)

    return sums

def sum_for_v6(packet):
    # hdr.ipv6.src_addr,               /* 128 */
    # hdr.ipv6.dst_addr,               /* 128 */
    # meta.length_without_ip_header,   /* 32 */
    # 24w0,                            /* 24 */
    # 		hdr.ipv6.next_header,            /* 8 */
    #                                  /* total: 324 */

    # 		// UDP header
    # hdr.udp.src_port,                /* 16 */
    # hdr.udp.dst_port,                /* 16 */
    # hdr.udp.payload_length           /* 16 */
    #                                  /* all: 372 */

    # order does not matter!
    sums = ""

    sums += ipaddress.IPv6Address(packet[IPv6].src.decode("utf-8")).packed
    sums += ipaddress.IPv6Address(packet[IPv6].dst.decode("utf-8")).packed
    sums += struct.pack("H", packet[IPv6].plen)
    sums += struct.pack("B", packet[IPv6].nh)

#    print("{} - {} - {}".format(len(sums), sums, checksum_scapy(sums)))

    return sums

def sum_for_v4(packet):
    # Get diffs -- for UDP
    # udp_v4 =
    #             hdr.ipv4.src_addr,
    #             hdr.ipv4.dst_addr,
    #             8w0,
    #             hdr.ipv4.protocol,
    #             meta.length_without_ip_header,
    sums = ""

    sums += ipaddress.IPv4Address(packet[IP].src.decode("utf-8")).packed
    sums += ipaddress.IPv4Address(packet[IP].dst.decode("utf-8")).packed
    sums += struct.pack("H", packet[IP].len - 20) # -20 for ip header
    sums += struct.pack("H", packet[IP].proto) # udp / tcp

    return sums

if __name__ == '__main__':
    p = []

    e0 = Ether(src="00:00:00:00:00:00",
               dst="00:00:00:00:00:00")

    i0 = IP(src = "0.0.0.0",
            dst = "0.0.0.0")

    t0 = TCP(dport=0, sport=0)
    #t = TCP(dport=80, sport=random.randint(49152,65535))

    #    print("chk_t = {}".format(t))

    e = Ether(src="02:53:55:42:45:01", dst='ff:ff:ff:ff:ff:ff')
    i4 = IP(src = "192.168.1.1", dst = "192.168.4.2")
    i6 = IPv6(src = "2001:db8:42::1", dst = "2001:db8::2")
    i62 = IPv6(src = "2001:db8:42::2", dst = "2001:db8::2")
    t = TCP(dport=80, sport=1337)
    u = UDP(dport=80, sport=1337)

    #print("chk_t = {}".format(t))

    d0 = ""
    d = "A"

    p6_udp = e / i6 / u / d
    p6_udp2 = e / i62 / u / d
    p4_udp = e / i4 / u / d

    p4_p6_1 = {
        "ipv6": p6_udp,
        "ipv4": p4_udp
    }
    p.append(p4_p6_1)

    p4_p6_2 = {
        "ipv6": p6_udp2,
        "ipv4": p4_udp
    }
    p.append(p4_p6_2)

    #p.append(e / i62 / u / d)



    for p_pair in p:
        v6 = p_pair["ipv6"]
        v4 = p_pair["ipv4"]

        checksums = {}
        header_checksums = {}
        diff_ip_headers = 0

        for packet in [v6, v4]:
            #print("p = {}".format(packet.__repr__()))
            packet_rebuild = packet.__class__(str(packet))
            print("rebuild = {}".format(packet_rebuild.__repr__()))
            chk_old = packet[UDP].chksum
            chk_new = packet_rebuild[UDP].chksum
            # print("chk1 = {} chk2={}".format(chk_old, chk_new))

            sums = ""


            if IPv6 in packet:
                headertype = "ipv6"
                sums += sum_for_v6(packet_rebuild)

                if UDP in packet:
                    checksums["ipv6"] = packet_rebuild[UDP].chksum

            if IP in packet:
                headertype = "ipv4"
                sums += sum_for_v4(packet_rebuild)

                if UDP in packet:
                    checksums["ipv4"] = packet_rebuild[UDP].chksum

            header_checksums[headertype] = checksum_scapy(sums)
            print("Checksum-parts {} for {}".format(checksum_scapy(sums), packet_rebuild.__repr__()))

        print("UDP v6: {} v4: {} diff: {}".format(checksums["ipv6"],
                                                  checksums["ipv4"],
                                                  checksums["ipv6"] - checksums["ipv4"]))

        header_diff_from_v6 = header_checksums["ipv6"] - header_checksums["ipv4"]
        print("Header v6: {} v4: {} diff: {}".format(header_checksums["ipv6"],
                                                     header_checksums["ipv4"],
                                                     header_diff_from_v6))

        print("Translating v6 to v4, expected v4 UDP checksum: {}".format(
            checksums["ipv6"] - header_diff_from_v6))

        print("Translating v4 to v6, expected v6 UDP checksum: {}".format(
            checksums["ipv4"] + header_diff_from_v6))