"""
Packet read and write routines for pcap format.
See http://wiki.wireshark.org/Development/LibpcapFileFormat
"""
import logging
import types

from packetracer import packetracer
from packetracer.structcbs import pack_I, unpack_I_le
from packetracer.structcbs import pack_IIII, unpack_IIII, unpack_IIII_le
from packetracer.layer12 import ethernet, linuxcc, radiotap, btle, can

logger = logging.getLogger("packetracer")


# PCAP/TCPDump related
# PCAP file header

# File magic numbers
# pcap using microseconds resolution
TCPDUMP_MAGIC	        	= 0xA1B2C3D4
TCPDUMP_MAGIC_SWAPPED	        = 0xD4C3B2A1
# pcap using nanoseconds resolution
TCPDUMP_MAGIC_NANO		= 0xA1B23C4D
TCPDUMP_MAGIC_NANO_SWAPPED	= 0x4D3CB2A1

PCAP_VERSION_MAJOR		= 2
PCAP_VERSION_MINOR		= 4

DLT_NULL				= 0
DLT_EN10MB				= 1
DLT_EN3MB				= 2
DLT_AX25				= 3
DLT_PRONET				= 4
DLT_CHAOS				= 5
DLT_IEEE802				= 6
DLT_ARCNET				= 7
DLT_SLIP				= 8
DLT_PPP					= 9
DLT_FDDI				= 10
DLT_PFSYNC				= 18
DLT_IEEE802_11				= 105
DLT_LINUX_SLL				= 113
DLT_PFLOG				= 117
DLT_IEEE802_11_RADIO			= 127
DLT_CAN_SOCKETCAN		        = 227
DLT_LINKTYPE_BLUETOOTH_LE_LL		= 251
LINKTYPE_BLUETOOTH_LE_LL_WITH_PHDR	= 256

PCAPTYPE_CLASS = {
	DLT_LINUX_SLL: linuxcc.LinuxCC,
	DLT_EN10MB: ethernet.Ethernet,
	DLT_CAN_SOCKETCAN: can.CAN,
	DLT_IEEE802_11_RADIO: radiotap.Radiotap,
	LINKTYPE_BLUETOOTH_LE_LL_WITH_PHDR: btle.BTLEHdr
}


class PcapFileHdr(packetracer.Packet):
	"""pcap file header."""
	# header length = 24
	__hdr__ = (
		("magic", "I", TCPDUMP_MAGIC),
		("v_major", "H", PCAP_VERSION_MAJOR),
		("v_minor", "H", PCAP_VERSION_MINOR),
		("thiszone", "I", 0),
		("sigfigs", "I", 0),
		("snaplen", "I", 1500),
		("linktype", "I", 1),
	)


class PcapPktHdr(packetracer.Packet):
	"""pcap packet header."""
	# header length: 16
	__hdr__ = (
		("tv_sec", "I", 0),
		# this can be either microseconds or nanoseconds: check magic number
		("tv_usec", "I", 0),
		("caplen", "I", 0),
		("len", "I", 0),
	)


# PCAP callbacks


def pcap_cb_init_write(self, snaplen=1500, linktype=DLT_EN10MB, **initdata):
	self._timestamp = 0
	header = PcapFileHdr(magic=TCPDUMP_MAGIC_NANO, snaplen=snaplen, linktype=linktype)
	self._fh.write(header.bin())


def pcap_cb_write(self, bts, **metadata):
	ts = metadata.get("ts", self._timestamp + 1000000)
	self._timestamp = ts
	sec = int(ts / 1000000000)
	nsec = ts - (sec * 1000000000)

	# logger.debug("paket time sec/nsec: %d/%d", sec, nsec)
	n = len(bts)
	self._fh.write(pack_IIII(sec, nsec, n, n))
	self._fh.write(bts)


def pcap_cb_init_read(self, **initdata):
	buf = self._fh.read(24)
	# file header is skipped per default (needed for __next__)
	self._fh.seek(24)
	# this is not needed anymore later on but we set it anyway
	fhdr = PcapFileHdr(buf)
	self._closed = False

	if fhdr.magic not in [TCPDUMP_MAGIC, TCPDUMP_MAGIC_NANO, TCPDUMP_MAGIC_SWAPPED, TCPDUMP_MAGIC_NANO_SWAPPED]:
		return False
	is_le = False

	# handle file types
	if fhdr.magic == TCPDUMP_MAGIC:
		self._resolution_factor = 1000
		# Note: we could use PcapPktHdr/PcapLEPktHdr to parse pre-packetdata but calling unpack directly
		# greatly improves performance
		self._callback_unpack_meta = unpack_IIII
	elif fhdr.magic == TCPDUMP_MAGIC_NANO:
		self._resolution_factor = 1
		self._callback_unpack_meta = unpack_IIII
	elif fhdr.magic == TCPDUMP_MAGIC_SWAPPED:
		is_le = True
		self._resolution_factor = 1000
		self._callback_unpack_meta = unpack_IIII_le
	elif fhdr.magic == TCPDUMP_MAGIC_NANO_SWAPPED:
		is_le = True
		self._resolution_factor = 1
		self._callback_unpack_meta = unpack_IIII_le
	else:
		raise ValueError("invalid tcpdump header, magic value: %s" % fhdr.magic)

	linktype = fhdr.linktype if not is_le else unpack_I_le(pack_I(fhdr.linktype))[0]
	self._lowest_layer_new = PCAPTYPE_CLASS.get(linktype, None)

	def is_resolution_nano(obj):
		"""return -- True if resolution is in Nanoseconds, False if milliseconds."""
		return obj._resolution_factor == 1000

	self.is_resolution_nano = types.MethodType(is_resolution_nano, self)
	return True


def pcap_cb_read(self):
	buf = self._fh.read(16)

	if not buf:
		raise StopIteration

	d = self._callback_unpack_meta(buf)
	buf = self._fh.read(d[2])

	return d[0] * 1000000000 + (d[1] * self._resolution_factor), buf


def pcap_cb_btstopkt(self, meta, bts):
	return self._lowest_layer_new(bts)


# PCAPNG related_btstopkt
# Generic/filetype invariant related

FILETYPE_PCAP	= 0
FILETYPE_PCAPNG	= 1

# type_id : [
#	cb_init_write(obj, **initdata),
#	cb_write(self, bytes, **metadata),
#	cb_init_read(obj, **initdata),
#	cb_read(self): metadata, bytes
#	cb_btstopkt(self, metadata, bytes): pkt
# ]
FILEHANDLER = {
	FILETYPE_PCAP: [
		pcap_cb_init_write, pcap_cb_write, pcap_cb_init_read, pcap_cb_read, pcap_cb_btstopkt
	]
	#FILETYPE_PCAPNG : [
	#	None, None, None
	#]
}


class FileHandler(object):
	def __init__(self, filename, accessmode):
		self._fh = open(filename, accessmode)
		self._closed = False

	def __enter__(self):
		return self

	def __exit__(self, objtype, value, traceback):
		self.close()

	def flush(self):
		self._fh.flush()

	def close(self):
		self._closed = True
		self._fh.close()


class PcapHandler(FileHandler):
	MODE_READ = 1
	MODE_WRITE = 2

	def __init__(self, filename, mode, filetype=FILETYPE_PCAP, **initdata):
		try:
			callbacks = FILEHANDLER[filetype]
		except IndexError:
			raise Exception("unknown filehandler type for mode %d: %d" % (mode, filetype))

		if mode == PcapHandler.MODE_WRITE:
			super().__init__(filename, "wb")
			callbacks[0](self, **initdata)
			self.write = types.MethodType(callbacks[1], self)
		elif mode == PcapHandler.MODE_READ:
			super().__init__(filename, "rb")
			ismatch = False

			for pcaptype, callbacks in FILEHANDLER.items():
				self._fh.seek(0)
				# init callback
				ismatch = callbacks[2](self, **initdata)

				if ismatch:
					#logger.debug("found handler for file: %x", pcaptype)
					# read callback
					self.__next__ = types.MethodType(callbacks[3], self)
					# bytes-to-packet callback
					self._btstopkt = types.MethodType(callbacks[4], self)
					break
			if not ismatch:
				raise Exception("no matching handler found")
		else:
			raise Exception("wrong mode: %d" % mode)

	def read_packet(self, pktfilter=lambda pkt: True):
		"""
		pktfilter -- filter as lambda function to match packets to be retrieved,
			return True to accept a specific packet.
		return -- (metadata, packet) if packet can be created from bytes
			else (metadata, bytes). For pcap/tcpdump metadata is a nanoseconds timestamp
		"""
		while True:
			# until StopIteration
			meta, bts = self.__next__()

			try:
				pkt = self._btstopkt(meta, bts)
			except Exception as ex:
				logger.warning("could not create packets from bytes: %r", ex)
				return meta, bts

			if pktfilter(pkt):
				return meta, pkt

	def read_packet_iter(self, pktfilter=lambda pkt: True):
		"""
		pktfilter -- filter as lambda function to match packets to be retrieved,
			return True to accept a specific packet.
		return -- iterator yielding (metadata, packet)
		"""
		if self._closed:
			return

		while True:
			try:
				yield self.read_packet(pktfilter=pktfilter)
			except:
				return

	def __iter__(self):
		"""
		return -- (metadata, bytes)
		"""
		if self._closed:
			return

		while True:
			try:
				yield self.__next__()
			except StopIteration:
				break


class Writer(PcapHandler):
	"""
	Simple pcap writer supporting pcap format.
	"""
	def __init__(self, filename, filetype=FILETYPE_PCAP, **initdata):
		super().__init__(filename, PcapHandler.MODE_WRITE, **initdata)


class Reader(PcapHandler):
	"""
	Simple pcap file reader supporting pcap format.
	"""
	def __init__(self, filename, filetype=FILETYPE_PCAP, **initdata):
		super().__init__(filename, PcapHandler.MODE_READ, **initdata)
