# -*- coding: utf-8 -*-
import logging
import sys

logging.basicConfig(level=logging.DEBUG)
sys.path.insert(0, "./netzob-sygus/src/")  # For Netzob

from netzob.all import *

# UDP symbol
udp_src = Field(name='sport', domain=Raw(nbBytes=2))
udp_dst = Field(name='dport', domain=Raw(nbBytes=2))
udp_length = Field(name='length', domain=Raw(nbBytes=2, unitSize=AbstractType.UNITSIZE_16))
udp_checksum = Field(name='checksum', domain=Raw(nbBytes=2, unitSize=AbstractType.UNITSIZE_16))
udp_payload = Field(name='payload', domain=Raw(b"\x61\x61\x61\x61\x61\x0a"))

udp_header = [udp_src, udp_dst, udp_length, udp_checksum, udp_payload]
udp_length.domain = Size(udp_header, dataType=Raw(nbBytes=2, unitSize=AbstractType.UNITSIZE_16), factor=1/float(8))

# Pseudo IP header
ip_pseudo_saddr = Field(name='IP Source address', domain=IPv4("192.168.0.1"))
ip_pseudo_daddr = Field(name='IP Destination address', domain=IPv4("192.168.0.2"))
ip_pseudo_length = Field(name='UDP Length', domain=Size(udp_header, dataType=Raw(nbBytes=2, unitSize=AbstractType.UNITSIZE_16), factor=1/float(8)))
ip_pseudo_zero = Field(name='Zero', domain=Raw(b"\x00"))
ip_pseudo_proto = Field(name='Protocol', domain=Integer(value=17, unitSize=AbstractType.UNITSIZE_8))

ip_pseudo_header = Field(name='Pseudo IP header', isPseudoField=True)
ip_pseudo_header.fields = [ip_pseudo_saddr, ip_pseudo_daddr, ip_pseudo_zero, ip_pseudo_proto, ip_pseudo_length]

udp_checksum.domain = InternetChecksum([ip_pseudo_header] + udp_header, dataType=Raw(nbBytes=2, unitSize=AbstractType.UNITSIZE_16))

udp_symbol = Symbol(name='UDP layer', fields=udp_header+[ip_pseudo_header])

print(udp_symbol._str_debug())
dot = udp_symbol.generateDot()
print(dot)

f = open("/tmp/plop.dot", 'w')
f.write(dot)
f.close

import subprocess
subprocess.Popen("xdot /tmp/plop.dot", shell=True)