1"""
2User Datagram Protocol (UDP)
3
4RFC 768 - User Datagram Protocol
5RFC 2460 - Internet Protocol, Version 6 (IPv6) Specification
6RFC 2675 - IPv6 Jumbograms
7RFC 4113 - Management Information Base for the UDP
8RFC 5405 - Unicast UDP Usage Guidelines for Application Designers
9"""
10import logging
11import struct
12
13from pypacker import pypacker, checksum
14from pypacker.pypacker import FIELD_FLAG_AUTOUPDATE, FIELD_FLAG_IS_TYPEFIELD
15# handler
16from pypacker.layer567 import telnet, tftp, dns, dhcp, iso15118, ntp, rtp, sip, pmap, radius, stun
17from pypacker.structcbs import unpack_H, pack_ipv4_header, pack_ipv6_header
18
19# avoid references for performance reasons
20in_cksum = checksum.in_cksum
21
22logger = logging.getLogger("pypacker")
23
24UDP_PORT_MAX	= 65535
25
26
27UDP_PROTO_TELNET	= 23
28UDP_PROTO_DNS		= (53, 5353)
29UDP_PROTO_DHCP		= (67, 68)
30UDP_PROTO_TFTP		= 69
31UDP_PROTO_PMAP		= 111
32UDP_PROTO_NTP		= 123
33UDP_PROTO_RADIUS	= (1812, 1813, 1645, 1646)
34UDP_PROTO_STUN		= 3478
35UDP_PROTO_RTP		= (5004, 5005)
36UDP_PROTO_SIP		= (5060, 5061)
37UDP_PROTO_ISO15118	= 15118
38
39
40class UDP(pypacker.Packet):
41	__hdr__ = (
42		("sport", "H", 0xDEAD),
43		("dport", "H", 0, FIELD_FLAG_AUTOUPDATE | FIELD_FLAG_IS_TYPEFIELD),
44		("ulen", "H", 8, FIELD_FLAG_AUTOUPDATE),  # header + body, min 8
45		("sum", "H", 0, FIELD_FLAG_AUTOUPDATE)
46	)
47
48	__handler__ = {
49		UDP_PROTO_TELNET: telnet.Telnet,
50		UDP_PROTO_TFTP: tftp.TFTP,
51		UDP_PROTO_DNS: dns.DNS,
52		UDP_PROTO_DHCP: dhcp.DHCP,
53		UDP_PROTO_ISO15118: iso15118.SDP,
54		UDP_PROTO_PMAP: pmap.Pmap,
55		UDP_PROTO_NTP: ntp.NTP,
56		UDP_PROTO_RADIUS: radius.Radius,
57		UDP_PROTO_RTP: rtp.RTP,
58		UDP_PROTO_SIP: sip.SIP,
59		UDP_PROTO_STUN: stun.STUN
60	}
61
62	def _update_fields(self):
63		# UDP-checksum needs to be updated on one of the following:
64		# - this layer itself or any upper layer changed
65		# - changes to the IP-pseudoheader
66		# There is no update on user-set checksums.
67		#changed = self._changed()
68		update = True
69
70		if self.ulen_au_active:
71			self.ulen = len(self)
72
73		#self._update_higherlayer_id()
74
75		try:
76			# changes to IP-layer, don't mind if this isn't IP
77			if not self._lower_layer._header_changed:
78				# lower layer doesn't need update, check for changes in present and upper layer
79				# logger.debug("lower layer did NOT change!")
80				update = True
81		except AttributeError:
82			# assume not an IP packet: we can't calculate the checksum
83			update = False
84
85		if update and self.sum_au_active:
86			self._calc_sum()
87
88	def _dissect(self, buf):
89		ports = [unpack_H(buf[0:2])[0], unpack_H(buf[2:4])[0]]
90
91		try:
92			# source or destination port should match
93			htype = [x for x in ports if x in pypacker.Packet._id_handlerclass_dct[UDP]][0]
94			self._init_handler(htype, buf[8:])
95		except:
96			# no type found
97			# logger.debug("could not parse type: %d because: %s" % (type, e))
98			pass
99		return 8
100
101	def _calc_sum(self):
102		"""Recalculate the UDP-checksum."""
103		# TCP and underwriting are freaky bitches: we need the IP pseudoheader to calculate their checksum
104		# logger.debug("UDP sum recalc, sport=%s/dport=%s" % (self.sport, self.dport))
105		try:
106			# we need src/dst for checksum-calculation
107			src, dst = self._lower_layer.src, self._lower_layer.dst
108			#logger.debug(src + b" / "+ dst)
109			self.sum = 0
110			udp_bin = self.header_bytes + self.body_bytes
111
112			# IP-pseudoheader, check if version 4 or 6
113			if len(src) == 4:
114				s = pack_ipv4_header(src, dst, 17, len(udp_bin))  # 17 = UDP
115			else:
116				s = pack_ipv6_header(src, dst, 17, len(udp_bin))  # 17 = UDP
117
118			csum = in_cksum(s + udp_bin)
119
120			if csum == 0:
121				csum = 0xFFFF    # RFC 768, p2
122
123			# get the checksum of concatenated pseudoheader+TCP packet
124			# assign via non-shadowed variable to trigger re-packing
125			self.sum = csum
126		except (AttributeError, struct.error):
127			# not an IP packet as lower layer (src, dst not present) or invalid src/dst
128			pass
129
130	def direction(self, other):
131		direction = 0
132		# logger.debug("checking direction: %s<->%s" % (self, other))
133		if self.sport == other.sport and self.dport == other.dport:
134			direction = pypacker.Packet.DIR_SAME
135		if self.sport == other.dport and self.dport == other.sport:
136			direction = pypacker.Packet.DIR_REV
137		if direction == 0:
138			return pypacker.Packet.DIR_UNKNOWN
139		return direction
140
141	def reverse_address(self):
142		self.sport, self.dport = self.dport, self.sport
143