#
# Copyright 2021 John-Mark Gurney.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.


# from: https://en.wikipedia.org/wiki/Reserved_IP_addresses
# TEST-NET-1: 192.0.2.0/24
# TEST-NET-2: 198.51.100.0/24
# TEST-NET-3: 203.0.113.0/24

from unittest.mock import patch
from mocks import *
import unittest

# Silence useless warning: WARNING: No IPv4 address found on
import scapy.error
scapy.error.warning = lambda *a, **kw: None
from scapy.all import *

from kvm import _TestCase as _KVMTestCase
from kvm import KVM

from bpf import Timeval
from ctypes import Structure, Union, POINTER, sizeof, create_string_buffer, byref
from ctypes import c_char, c_void_p, c_uint8, c_uint16, c_uint32, c_int64, c_uint64

import asyncio
import fcntl
import functools
import itertools
import re
import sockio
import string
import sys

class if_data(Structure):
	_fields_ = [
		('ifi_type', c_uint8),
		('ifi_physical', c_uint8),
		('ifi_addrlen', c_uint8),
		('ifi_hdrlen', c_uint8),
		('ifi_link_state', c_uint8),
		('ifi_vhid', c_uint8),
		('ifi_datalen', c_uint16),
		('ifi_mtu', c_uint32),
		('ifi_metric', c_uint32),
		('ifi_baudrate', c_uint64),
		('ifi_ipackets', c_uint64),
		('ifi_ierrors', c_uint64),
		('ifi_opackets', c_uint64),
		('ifi_collisions', c_uint64),
		('ifi_ibytes', c_uint64),
		('ifi_obytes', c_uint64),
		('ifi_imcasts', c_uint64),
		('ifi_omcasts', c_uint64),
		('ifi_iqdrops', c_uint64),
		('ifi_oqdrops', c_uint64),
		('ifi_noproto', c_uint64),
		('ifi_hwassist', c_uint64),
		('ifi_epoch', c_int64), # XXX - broken on i386
		('ifi_lastchange', Timeval), # XXX - broken on 32-bit platforms
	]

class ifr_ifru(Union):
	_fields_ = [
		('ifru_data', POINTER(c_char)),
	]

class ifreq(Structure):
	_fields_ = [
		('ifr_name', c_char * 16),
		('ifr_ifru', ifr_ifru),
	]

def _makeflags(s):
	return set(s.split(','))

def _parsemedia(s):
	reg = '^Ethernet autoselect (1000baseT <full-duplex>)$'
	m = re.match(reg, s)

	return dict(ethernet='autoselect', media=dict(medium='1000baseT', options={ 'full-duplex' }))

__ifreqsock = None
def get_ifreqsock():
	global __ifreqsock

	if __ifreqsock == None:
		__ifreqsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

	return __ifreqsock

__kvm = None
def get_kvm():
	global __kvm

	if __kvm == None:
		__kvm = KVM()

	return __kvm

def if_data(iface):
	s = get_ifreqsock()

	kd = get_kvm()

	ifdatalen = kd.structsize('struct if_data')
	ifdata = create_string_buffer(ifdatalen)

	ifq = ifreq()

	ifq.ifr_name = iface.encode('us-ascii')
	ifq.ifr_ifru.ifru_data = ifdata

	r = fcntl.ioctl(s, sockio.SIOCGIFDATA, ifq)

	if r != 0:
		raise RuntimeError('ioctl returned %d')

	return kd.getstruct('struct if_data', ifdata)

async def waitcarrier(iface):
	while True:
		res = await ifconfig(iface)

		if res[iface]['status'] == 'active':
			return

		await asyncio.sleep(.5)

async def ifconfig(iface, *args, **kwargs):
	preargs = ()
	args = sum([ [k, str(v)] for k, v in kwargs.items() ], list(args))

	# Get detailed info about the interface if nothing is specified
	if not args:
		preargs = ('-m',)

	proc = await asyncio.create_subprocess_exec('/sbin/ifconfig', *preargs,
		iface, *args, stdout=asyncio.subprocess.PIPE,
		stderr=asyncio.subprocess.PIPE)

	stdout, stderr = await proc.communicate()

	if proc.returncode != 0:
		raise RuntimeError(stderr.decode('us-ascii').strip())

	if not stdout:
		return None

	# got something, parse it
	#print('ifc:', repr(stdout))
	stdout = stdout.decode('us-ascii')
	lines = stdout.split('\n')

	reg = '^(?P<iface>^.+): flags=[0-9a-f]{4,4}<(?P<flags>([A-Z0-9_]+(,[A-Z0-9_]+)*)?)> metric (?P<metric>[0-9]+) mtu (?P<mtu>[0-9]+)$'
	m = re.match(reg, lines[0])

	iface = m.group('iface')
	res = { iface: dict(mtu=int(m.group('mtu')), metric=int(m.group('metric')), flags=_makeflags(m.group('flags'))) }

	obj = res[iface]

	for i in lines[1:]:
		reg = '^\\toptions=[0-9a-f]{5,5}<(?P<options>([A-Z0-9_]+(,[A-Z0-9_]+)*)?)>$'
		m = re.match(reg, i)
		if m:
			obj['options'] = _makeflags(m.group('options'))
			continue

		reg = '^\\tcapabilities=[0-9a-f]{5,6}<(?P<capabilities>([A-Z0-9_]+(,[A-Z0-9_]+)*)?)>$'
		m = re.match(reg, i)
		if m:
			obj['capabilities'] = _makeflags(m.group('capabilities'))
			continue

		reg = '^\\tether (?P<ether>[0-9a-f]{2,2}(:[0-9a-f]{2,2}){5,5})$'
		m = re.match(reg, i)
		if m:
			obj['ether'] = m.group('ether')
			continue

		reg = '^\\tmedia: (?P<media>.+)$'
		m = re.match(reg, i)
		if m:
			obj['media'] = _parsemedia(m.group('media'))
			continue

		reg = '^\\tstatus: (?P<status>.+)$'
		m = re.match(reg, i)
		if m:
			obj['status'] = m.group('status')
			continue

		reg = '^\\tnd6 options=[0-9a-f]{2,2}<(?P<nd6options>([A-Z0-9_]+(,[A-Z0-9_]+)*)?)>$'
		m = re.match(reg, i)
		if m:
			obj['nd6'] = _makeflags(m.group('nd6options'))
			continue

	return res

async def asendp(*args, **kwargs):
	'''a coroutine wrapping scapy sendp.  All the arguments are
	the same, but you must await to get the results.'''

	loop = asyncio.get_running_loop()

	return await loop.run_in_executor(None, functools.partial(sendp, *args, **kwargs))

class PacketManager(object):
	def __init__(self, iface):
		'''Create a context manager for sniffing packets on the
		interface specified by iface.

		Sniffing will only begin once used as a context manager.

		Example:
		async with PacketManager(iface) as pm:
			pkt = await pm.getpkt()
		'''

		self._iface = iface
		self._sniffer = None

	def enqueuepkt(self, pkt):
		'''Internal function to enqueue a received packet.'''

		# hopefully these routines are executed in order they are
		# scheduled.  The Python docs does not assure that this this
		# will happen.
		#print('enq:', repr(pkt))
		asyncio.run_coroutine_threadsafe(self._queue.put(pkt), self._loop)

	async def __aenter__(self):
		if self._sniffer is not None:
			raise RuntimeError('already in a context')

		self._loop = asyncio.get_running_loop()
		self._queue = asyncio.Queue()
		self._sniffer = AsyncSniffer(iface=self._iface, prn=self.enqueuepkt)
		self._sniffer.start()

		return self

	async def __aexit__(self, exc_type, exc_value, traceback):
		await self._loop.run_in_executor(None, self._sniffer.stop)

	def getpkt(self, timeout=None):
		'''coroutine.  Return the next available packet.  If timeout is
		specified (in seconds), if a packet is not available in the timeout
		specified, the exception asyncio.TimeoutError is raised.'''

		return asyncio.wait_for(self._queue.get(), timeout)

_mkpktdata = lambda: itertools.cycle(string.printable)
_getpktdata = lambda *args: ''.join(itertools.islice(_mkpktdata(), *args)).encode()

# Convert capability flags to their respective ifconfig command argument
flagtoopt = dict(
	TXCSUM_IPV6='txcsum6',
	RXCSUM_IPV6='rxcsum6',
	TXCSUM='txcsum',
	RXCSUM='rxcsum',
	VLAN_HWTAGGING='vlanhwtag',
)

async def csuminternal(flag, testiface, checkiface):
	# XXX - I can't figure out how to get checksum to ensure that
	# the hardcoded check will ALWAYS be invalid
	if flag[0] == 'T':
		txiface, rxiface = testiface, checkiface
	else:
		return
		txiface, rxiface = checkiface, testiface

	# base packet
	p = Ether(src=get_if_hwaddr(txiface), dst=get_if_hwaddr(rxiface))

	# https://en.wikipedia.org/wiki/Reserved_IP_addresses
	if flag.endswith('_IPV6'):
		p = p / IP(src='192.168.0.1', dst='192.168.0.2',
		    chksum=0xbadc, flags='DF')
	else:
		p = p / IPv6(src='fc00:0b5d:041c:7e37::7e37',
		    dst='fc00:0b5d:041c:7e37::c43c')

	tcp = p / TCP(dport=443, flags='S', chksum=0xbadc)
	udp = p / UDP(dport=53, chksum=0xbadc) / \
	    b'this is a udp checksum test packet'


	await ifconfig(checkiface, '-' + flagtoopt[flag])

	for pref in [ '-', '' ]:
		await ifconfig(testiface, pref + flagtoopt[flag])

		await ifconfig(testiface, 'up')
		await ifconfig(checkiface, 'up')

		await waitcarrier(testiface)
		await waitcarrier(checkiface)

		async with PacketManager(rxiface) as pm:
			await asendp(tcp, iface=txiface, verbose=False)

			try:
				rpkt = await pm.getpkt(timeout=.5)
			except asyncio.TimeoutError:
				raise RuntimeError('failed to receive checksum test')

			print('recv:', repr(rpkt))


async def csumtest(testiface, checkiface):
	if False:
		raise ValueError('cannot be implemented this way, need the host'
		    'stack to set special mbuf flags, etc')
	csumtests = { 'RXCSUM', 'TXCSUM', 'RXCSUM_IPV6', 'TXCSUM_IPV6' }
	ifc = await ifconfig(testiface)

	print(ifc)
	for csum in csumtests.intersection(ifc[testiface]['capabilities']):
		print(repr(csum))
		await csuminternal(csum, testiface, checkiface)

async def mtucheck(sndiface, rcviface, mtusize):

	# make sure packet is padded out to full frame size
	p = Ether(src=get_if_hwaddr(sndiface), dst=get_if_hwaddr(rcviface)) / \
	    _getpktdata(mtusize - 14)

	#print('pktlen:', repr(p.build()))
	#print('sndiface:', repr(sndiface))
	#print('rcviface:', repr(rcviface))
	async with PacketManager(rcviface) as pm:
		try:
			await asendp(p, iface=sndiface, verbose=False)
		except OSError:
			raise RuntimeError('failed to send mtu size %d' %
			    mtusize)

		try:
			rpkt = await pm.getpkt(timeout=.5)
		except asyncio.TimeoutError:
			raise RuntimeError('failed to receive mtu size %d' %
			    mtusize)

		#print('got pkt:', repr(rpkt))
		if rpkt != p:
			raise RuntimeError(
			    'received packet did not match sent: %s != %s' %
			    (repr(rpkt), repr(p)))

async def mtutest(testiface, checkiface):
	for mtusize in [ 512, 1500, 1504, 2025, 4074, 7000, 8000, 9000, 9216 ]:
		try:
			await ifconfig(testiface, mtu=mtusize)
			await ifconfig(checkiface, mtu=mtusize)
		except RuntimeError:
			print('failed to set mtu %d, skipping...' % mtusize)
			continue

		print('mtu size:', mtusize)

		await ifconfig(testiface, 'up')
		await ifconfig(checkiface, 'up')

		await waitcarrier(testiface)
		await waitcarrier(checkiface)

		# if other way around ure won't work
		await mtucheck(testiface, checkiface, mtusize)
		await mtucheck(checkiface, testiface, mtusize)

async def hwassisttest(testiface):
	'''Check to make sure that the hwassist flags follow the
	ifconfig options.
	'''

	ifconf = await ifconfig(testiface)
	caps = ifconf[testiface]['capabilities']

	basecsumflags = { 'RXCSUM', 'TXCSUM', 'RXCSUM_IPV6', 'TXCSUM_IPV6' }

	# get the supported flags
	supcsumflags = caps & basecsumflags

	# turn everything off
	await ifconfig(testiface, *('-%s' % flagtoopt[x] for x in supcsumflags))

	# make sure it is off
	ifconf = await ifconfig(testiface)

	if ifconf[testiface]['flags'] & basecsumflags:
		raise RuntimeError('failed to clear all the csum flags.')

	# make sure _hwassist is 0
	raise NotImplementedError('tbd')

def main():
	testiface, checkiface = sys.argv[1:3]

	print('running...')

	if sys.argv[3] == 'mtu':
		asyncio.run(mtutest(testiface, checkiface))
	elif sys.argv[3] == 'csum':
		asyncio.run(csumtest(testiface, checkiface))
	elif sys.argv[3] == 'hwassist':
		asyncio.run(hwassisttest(testiface))

	print('done')

if __name__ == '__main__':
	main()

class _TestCase(unittest.IsolatedAsyncioTestCase):
	@patch('asyncio.sleep')
	@patch(__name__ + '.ifconfig')
	async def test_waitcarrier(self, ifc, asleep):
		cnt = [ 0 ]
		async def se(iface):
			cnt[0] += 1

			if cnt[0] < 5:
				return { iface: dict(status='no carrier') }

			return { iface: dict(status='active') }

		ifc.side_effect = se

		await waitcarrier('foo')

		ifc.assert_called_with('foo')

		# Better to have an event to wait on, but that isn't available
		asleep.assert_called_with(.5)

	@patch('asyncio.create_subprocess_exec')
	async def test_ifconf_err(self, cse):
		errmsg =  \
		    b'ifconfig: ioctl SIOCSIFMTU (set mtu): Invalid argument'
		wrap_subprocess_exec(cse, stderr=errmsg, retcode=1)

		with self.assertRaisesRegex(RuntimeError, re.escape(errmsg.decode())):
			await ifconfig('foobar', mtu=7000)

	@patch('asyncio.create_subprocess_exec')
	async def test_ifconf(self, cse):
		wrap_subprocess_exec(cse, b'')

		self.assertIsNone(await ifconfig('foobar', mtu=4000))

		cse.assert_called_with('/sbin/ifconfig', 'foobar',
		    'mtu', '4000', stdout=asyncio.subprocess.PIPE,
		    stderr=asyncio.subprocess.PIPE)

		wrap_subprocess_exec(cse, b'')

		self.assertIsNone(await ifconfig('foobar', 'up'))

		cse.assert_called_with('/sbin/ifconfig', 'foobar',
		    'up', stdout=asyncio.subprocess.PIPE,
		    stderr=asyncio.subprocess.PIPE)

		output = b'''foobar: flags=8843<UP,BROADCAST,RUNNING,SIMPLEX,MULTICAST> metric 0 mtu 1500
	options=80188<VLAN_MTU,VLAN_HWCSUM,TSO4,LINKSTATE>
	capabilities=c019b<RXCSUM,TXCSUM,VLAN_MTU,VLAN_HWTAGGING,VLAN_HWCSUM,TSO4,VLAN_HWTSO,LINKSTATE>
	ether 52:12:34:56:78:90
	media: Ethernet autoselect (1000baseT <full-duplex>)
	status: active
	nd6 options=29<PERFORMNUD,IFDISABLED,AUTO_LINKLOCAL>
'''

		res = dict(foobar=dict(
		        flags={ 'UP','BROADCAST', 'RUNNING',
		        'SIMPLEX', 'MULTICAST' },
		        metric=0, mtu=1500,
			capabilities={ 'RXCSUM', 'TXCSUM', 'VLAN_MTU',
		            'VLAN_HWTAGGING', 'VLAN_HWCSUM', 'TSO4',
		            'VLAN_HWTSO', 'LINKSTATE' },
		        options={ 'VLAN_MTU', 'VLAN_HWCSUM', 'TSO4',
		            'LINKSTATE' },
		        ether='52:12:34:56:78:90',
		        media=dict(ethernet='autoselect',
		            media=dict(medium='1000baseT',
		                options={ 'full-duplex' })),
		        status='active',
		        nd6={ 'PERFORMNUD', 'IFDISABLED', 'AUTO_LINKLOCAL' },
		    )
		)

		wrap_subprocess_exec(cse, output)

		ret = await ifconfig('foobar')

		cse.assert_called_with('/sbin/ifconfig', '-m', 'foobar',
		    stdout=asyncio.subprocess.PIPE,
		    stderr=asyncio.subprocess.PIPE)

		self.assertEqual(ret, res)

		output = b'''foobar: flags=8843<UP,BROADCAST,RUNNING,SIMPLEX,MULTICAST> metric 0 mtu 1500
	options=80188<VLAN_MTU,VLAN_HWCSUM,TSO4,LINKSTATE>
	capabilities=68009b<RXCSUM,TXCSUM,VLAN_MTU,VLAN_HWTAGGING,VLAN_HWCSUM,LINKSTATE,RXCSUM_IPV6,TXCSUM_IPV6>
	ether 52:12:34:56:78:90
	media: Ethernet autoselect (1000baseT <full-duplex>)
	status: active
	nd6 options=29<PERFORMNUD,IFDISABLED,AUTO_LINKLOCAL>
'''

		res = dict(foobar=dict(
		        flags={ 'UP','BROADCAST', 'RUNNING',
		        'SIMPLEX', 'MULTICAST' },
		        metric=0, mtu=1500,
			capabilities={ 'RXCSUM', 'TXCSUM', 'VLAN_MTU',
		            'VLAN_HWTAGGING', 'VLAN_HWCSUM', 'LINKSTATE',
		            'RXCSUM_IPV6', 'TXCSUM_IPV6', },
		        options={ 'VLAN_MTU', 'VLAN_HWCSUM', 'TSO4',
		            'LINKSTATE' },
		        ether='52:12:34:56:78:90',
		        media=dict(ethernet='autoselect',
		            media=dict(medium='1000baseT',
		                options={ 'full-duplex' })),
		        status='active',
		        nd6={ 'PERFORMNUD', 'IFDISABLED', 'AUTO_LINKLOCAL' },
		    )
		)

		wrap_subprocess_exec(cse, output)

		ret = await ifconfig('foobar')

		cse.assert_called_with('/sbin/ifconfig', '-m', 'foobar',
		    stdout=asyncio.subprocess.PIPE,
		    stderr=asyncio.subprocess.PIPE)

		self.assertEqual(ret, res)

		output = b'''foobar: flags=8843<UP,BROADCAST,RUNNING,SIMPLEX,MULTICAST> metric 0 mtu 1500
	options=80188<VLAN_MTU,VLAN_HWCSUM,TSO4,LINKSTATE>
	ether 52:12:34:56:78:90
	media: Ethernet autoselect (1000baseT <full-duplex>)
	status: no carrier
	nd6 options=29<PERFORMNUD,IFDISABLED,AUTO_LINKLOCAL>
'''

		res = dict(foobar=dict(
		        flags={ 'UP','BROADCAST', 'RUNNING',
		        'SIMPLEX', 'MULTICAST' },
		        metric=0, mtu=1500,
		        options={ 'VLAN_MTU', 'VLAN_HWCSUM', 'TSO4',
		            'LINKSTATE' },
		        ether='52:12:34:56:78:90',
		        media=dict(ethernet='autoselect',
		            media=dict(medium='1000baseT',
		                options={ 'full-duplex' })),
		        status='no carrier',
		        nd6={ 'PERFORMNUD', 'IFDISABLED', 'AUTO_LINKLOCAL' },
		    )
		)

		wrap_subprocess_exec(cse, output)

		ret = await ifconfig('foobar')

		self.assertEqual(ret, res)

	@patch(__name__ + '.sendp')
	async def test_asendp(self, sndp):
		kwargs = dict(a=5, b='foo')
		pkt = object()
		retv = object()

		sndp.return_value = retv

		r = await asendp(pkt, **kwargs)

		self.assertIs(r, retv)

		sndp.assert_called_with(pkt, **kwargs)

	@patch(__name__ + '.AsyncSniffer')
	async def test_pktmgr(self, asyncs):
		# That it can be a context manager
		async with PacketManager('iface') as pm:
			# that a RuntimeError is raised
			with self.assertRaises(RuntimeError):
				# when it's used as a context manager again
				async with pm as foo:
					pass

			# That it created a queue
			self.assertIsInstance(pm._queue, asyncio.Queue)

			# that it called asyncsniffer with the correct arguments
			asyncs.assert_called_with(iface='iface', prn=pm.enqueuepkt)

			# that it was started
			asyncs().start.assert_called()

			# that when a pkt
			pkt = object()

			# is enqueued
			pm.enqueuepkt(pkt)

			# that getpkt returns it
			self.assertIs(await pm.getpkt(), pkt)

			# that a timeout parameter can be provided
			with self.assertRaises(asyncio.TimeoutError):
				await pm.getpkt(timeout=0)

		# that when the context manager stops, stop is called
		asyncs().stop.assert_called()

	def test_ifdata(self):
		r = if_data('ue0')
		print('if:', repr(r))