import socket
import struct

def _checksum_func(data: bytes) -> int:
    """
    _checksum_func is an internal method to perform 1s compliment checksum calculation for TCP/UDP headers

    :rtype: int
    :param data: A byte object which represents the header for checksum calculation
    :return: checksum: an int representing the checksum of the bytes/header.
    """
    # https://github.com/houluy/UDP/blob/master/udp.py#L120
    checksum = 0
    data_len = len(data)
    if data_len % 2:
        data_len += 1
        data += struct.pack("!B", 0)
    for i in range(0, data_len, 2):
        w = (data[i] << 8) + (data[i + 1])
        checksum += w
    checksum = (checksum >> 16) + (checksum & 0xFFFF)
    checksum = ~checksum & 0xFFFF
    return checksum



def chksum(packet: bytes) -> int:
    s = 0       # Binary Sum
    # loop taking 2 characters at a time
    for i in range(0, len(packet), 2):
        if (i+1) < len(packet):
            a = packet[i]
            b = packet[i+1]
            s = s + (a+(b << 8))
        elif (i+1)==len(packet):
            s += packet[i]
        else:
            raise "Something Wrong here"
    # One's Complement
    s = s + (s >> 16)
    s = ~s & 0xffff
    return s


# IPv4 Header
ip_ver = 4
ip_ihl = 5
ttl = 5
ip_tos = 0
ip_tot_len = 0 ## Reset this once UDP header is written.
ip_id = 257
ip_frag_off = 0
l4_proto = 17
ip_check = 0
ip_saddr = socket.inet_aton("192.168.0.80")
ip_daddr = socket.inet_aton(socket.gethostbyname("1.1.1.1"))

ip_ihl_ver = (ip_ver << 4) + ip_ihl

ip_header = struct.pack(
    "!BBHHHBBH4s4s",
    ip_ihl_ver,
    ip_tos,
    ip_tot_len,
    ip_id,
    ip_frag_off,
    ttl,
    l4_proto,
    ip_check,
    ip_saddr,
    ip_daddr,
)



udp_src_port = 34500
udp_dst_port = 34500

data = "01234567".encode()
# UDP is a bit stupid, and takes a lower layer info as part of it's checksum. Specifically src/dst IP addr.
# This is called the pseudo header
pseudo_header = struct.pack("!BBH", 0, socket.getprotobyname("udp"), len(data) + 8 )
pseudo_header = ip_saddr + ip_daddr + pseudo_header
# Set the checksum to 0, so we can generate a header, then calculate the checksum and re-apply
checksum = 0
udp_header = struct.pack("!4H", udp_src_port, udp_dst_port, len(data) + 8 , checksum)
checksum = _checksum_func(pseudo_header + udp_header + data)
udp_header = struct.pack("!4H", udp_src_port, udp_dst_port, len(data) + 8 , checksum)


ip_tot_len = len(ip_header) + len(udp_header)

ip_header = struct.pack(
    "!BBHHHBBH4s4s",
    ip_ihl_ver,
    ip_tos,
    ip_tot_len,
    ip_id,
    ip_frag_off,
    ttl,
    l4_proto,
    ip_check,
    ip_saddr,
    ip_daddr,
)

ip_check = chksum(ip_header)

ip_header = struct.pack(
    "!BBHHHBBH4s4s",
    ip_ihl_ver,
    ip_tos,
    ip_tot_len,
    ip_id,
    ip_frag_off,
    ttl,
    l4_proto,
    ip_check,
    ip_saddr,
    ip_daddr,
)



raw_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_RAW)
raw_sock.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, 1)
raw_sock.sendto(ip_header + udp_header + data, ("1.1.1.1", 0))