#!/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


# 1. convert to array of "bytes"
# 2. import into an array
# 3. sum everything up


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"

    p.append(e / i4 / u / d)
    p.append(e / i6 / u / d)
    p.append(e / i62 / u / d)

    for packet in p:
        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 IP in packet:
            sums += sum_for_v4(packet_rebuild)
        if IPv6 in packet:
            sums += sum_for_v6(packet_rebuild)
#        if UDP in packet:
#            sums += sum_for_udp(packet_rebuild)

        print("Checksum-parts {} for {}".format(checksum_scapy(sums), packet_rebuild.__repr__()))


# Checksums:
# - tcp
# - udp
# - icmp6
# - icmp
# - ipv4 (no payload)
#
#         t.chksum = None