#!/usr/bin/env python3

# Code copied from /usr/lib/python3/dist-packages/scapy/utils.py

import struct
import array

if struct.pack("H",1) == b"\x00\x01": # big endian
    def checksum(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 & 0xffff
else:
    def checksum(pkt):
        # padding to even length
        if len(pkt) % 2 == 1:
            pkt += b"\0"

        # create array with unsigned shorts,
        # sum it up -> results in 16 bit uint (?)
        s = sum(array.array("H", pkt))
        print("sum={}".format(s))
        # This step is understood

        # (s & 0xffff) -> mask for 16 bit, truncate rest
        # (s >> 16) right shift by 16 bits
        # >>> 2**17 >> 16
        # 2
        # Basically:
        # add higher 16 bits to lower 16 bits

        s = (s >> 16) + (s & 0xffff)

        # Now we add the remaining possible 16 bits
        # i.e. above we take bits 0..15 + 16..31
        # now we add [0..15 + 16..31] + "32..47"
        # I assume this is for catching the carry
        # over bits from the previous calculation
        s += s >> 16

        # according to https://stackoverflow.com/questions/8305199/the-tilde-operator-in-python
        # ~ is the bitwise complement operator in python
        # which essentially calculates -x - 1

        # 2 complement
        s = ~s

        # (s>>8)&0xff) -> filter bits 8..15
        # | = bitwise OR
        # s<<8 : left shift everything in original s
        # | s<<8: ?
        # & 0xffff: filter / allow only 16 bits
        # Seems like the idea is to filter correctly for
        # 16 bit - unclear why "s & 0xffff" does not do the job

        return (((s>>8)&0xff)|s<<8) & 0xffff


if __name__ == '__main__':
    import sys
    pkt = sys.argv[1]
    pkt_bytes = pkt.encode('utf-8')
    print(pkt_bytes)
    print(checksum(pkt_bytes))

    ipv4 = "AAAA"
    ipv6 = "BBAA"