import socket
import struct

class Ipv4Header():

    def __init__(self, sourceIpAddress, destinationIpAddress):
        self.version = 4 # ipv4
        self.internetHeaderLength = 5 # number of 32 bit words
        self.differntiatedServicesCodePoint = 0
        self.electronicCommunicationNetwork = 0
        self.identification = 0 # only one packet will be sent
        self.flags = 0 # no flags needed
        self.fragmentOffset = 0 # packet won't be fragmented
        self.timeToLive = 128
        self.protocol = 6 # tcp
        self.headerChecksum = 0 # temporary number for checksum calculation
        self.sourceIpAddress = socket.inet_aton(sourceIpAddress) # format source ip
        self.destinationIpAddress = socket.inet_aton(destinationIpAddress) # format destination ip

    def calculateTotalLength(self, payloadLength):
        self.totalLength = payloadLength + (self.internetHeaderLength * 4)

    def packHeader(self):
        self.raw = struct.pack("!BBHHHBBH4s4s",
            (self.version << 4) + self.internetHeaderLength,
            (self.differntiatedServicesCodePoint << 2) + self.electronicCommunicationNetwork,
            self.totalLength,
            self.identification,
            (self.flags << 13) + self.fragmentOffset,
            self.timeToLive,
            self.protocol,
            self.headerChecksum,
            self.sourceIpAddress,
            self.destinationIpAddress)

    def calculateHeaderChecksum(self):
        s = 0
        w = 0
        for i in range(0, len(self.raw), 2):
            w = (self.raw[i] << 8) + self.raw[i + 1]
            s += w
        s = (s + (s >> 16)) & 0xffff
        s = ~s & 0xffff
        self.headerChecksum = s

if __name__ == "__main__":
    Ipv4Header = Ipv4Header("192.168.0.1","192.168.0.128")
    Ipv4Header.calculateTotalLength(0)
    Ipv4Header.packHeader()
    Ipv4Header.calculateHeaderChecksum()
    Ipv4Header.packHeader()

    print(Ipv4Header.raw)