lora-irrigation/lora.py

import asyncio
import functools
import os
import unittest

from Strobe.Strobe import Strobe
from Strobe.Strobe import AuthenticationFailed

domain = b'com.funkthat.lora.irrigation.shared.v0.0.1'

# Response to command will be the CMD and any arguments if needed.
# The command is encoded as an unsigned byte
CMD_TERMINATE = 1	# no args: terminate the sesssion, reply confirms

# The follow commands are queue up, but will be acknoledged when queued
CMD_WAITFOR = 2		# arg: (length): waits for length seconds
CMD_RUNFOR = 3		# arg: (chan, length): turns on chan for length seconds

class LORANode(object):
	'''Implement a LORANode initiator.'''

	def __init__(self, syncdatagram):
		self.sd = syncdatagram
		self.st = Strobe(domain)

	async def start(self):
		msg = self.st.send_enc(b'reqreset' + os.urandom(16)) + \
		    self.st.send_mac(8)

		resp = await self.sd.sendtillrecv(msg, 1)

		self.st.recv_enc(resp[:16])
		self.st.recv_mac(resp[16:])

		resp = await self.sd.sendtillrecv(
		    self.st.send_enc(b'confirm') + self.st.send_mac(8), 1)

		pkt = self.st.recv_enc(resp[:9])
		self.st.recv_mac(resp[9:])

		if pkt != b'confirmed':
			raise RuntimeError

	@staticmethod
	def _encodeargs(*args):
		r = []
		for i in args:
			r.append(i.to_bytes(4, byteorder='little'))

		return b''.join(r)

	async def _sendcmd(self, cmd, *args):
		cmdbyte = cmd.to_bytes(1, byteorder='little')
		pkt = await self.sd.sendtillrecv(
		    self.st.send_enc(cmdbyte +
		    self._encodeargs(*args)) + self.st.send_mac(8), 1)

		resp = self.st.recv_enc(pkt[:-8])
		self.st.recv_mac(pkt[-8:])

		if resp[0:1] != cmdbyte:
			raise RuntimeError('response does not match, got: %s, expected: %s' % (repr(resp[0:1]), repr(cmdbyte)))

	async def waitfor(self, length):
		return await self._sendcmd(CMD_WAITFOR, length)

	async def runfor(self, chan, length):
		return await self._sendcmd(CMD_RUNFOR, chan, length)

	async def terminate(self):
		return await self._sendcmd(CMD_TERMINATE)

class SyncDatagram(object):
	'''Base interface for a more simple synchronous interface.'''

	def __init__(self): #pragma: no cover
		pass

	async def recv(self, timeout=None): #pragma: no cover
		'''Receive a datagram.  If timeout is not None, wait that many
		seconds, and if nothing is received in that time, raise an TimeoutError
		exception.'''

		raise NotImplementedError

	async def send(self, data): #pragma: no cover
		'''Send a datagram.'''

		raise NotImplementedError

	async def sendtillrecv(self, data, freq):
		'''Send the datagram in data, every freq seconds until a datagram
		is received.  If timeout seconds happen w/o receiving a datagram,
		then raise an TimeoutError exception.'''

		while True:
			await self.send(data)
			try:
				return await self.recv(freq)
			except TimeoutError:
				pass

class MockSyncDatagram(SyncDatagram):
	'''A testing version of SyncDatagram.  Define a method runner which
	implements part of the sequence.  In the function, await on either
	self.get, to wait for the other side to send something, or await
	self.put w/ data to send.'''

	def __init__(self):
		self.sendq = asyncio.Queue()
		self.recvq = asyncio.Queue()
		self.task = None
		self.task = asyncio.create_task(self.runner())

		self.get = self.sendq.get
		self.put = self.recvq.put

	async def drain(self):
		'''Wait for the runner thread to finish up.'''

		return await self.task

	async def runner(self): #pragma: no cover
		raise NotImplementedError

	async def recv(self, timeout=None):
		return await self.recvq.get()

	async def send(self, data):
		return await self.sendq.put(data)

	def __del__(self): #pragma: no cover
		if self.task is not None and not self.task.done():
			self.task.cancel()

class TestSyncData(unittest.IsolatedAsyncioTestCase):
	async def test_syncsendtillrecv(self):
		class MySync(SyncDatagram):
			def __init__(self):
				self.sendq = []
				self.resp = [ TimeoutError(), b'a' ]

			async def recv(self, timeout=None):
				assert timeout == 1
				r = self.resp.pop(0)
				if isinstance(r, Exception):
					raise r

				return r

			async def send(self, data):
				self.sendq.append(data)

		ms = MySync()

		r = await ms.sendtillrecv(b'foo', 1)

		self.assertEqual(r, b'a')
		self.assertEqual(ms.sendq, [ b'foo', b'foo' ])

def timeout(timeout):
	def timeout_wrapper(fun):
		@functools.wraps(fun)
		async def wrapper(*args, **kwargs):
			return await asyncio.wait_for(fun(*args, **kwargs),
			    timeout)

		return wrapper

	return timeout_wrapper

class TestLORANode(unittest.IsolatedAsyncioTestCase):
	@timeout(2)
	async def test_lora(self):
		class TestSD(MockSyncDatagram):
			async def runner(self):
				l = Strobe(domain)

				# start handshake
				r = await self.get()

				pkt = l.recv_enc(r[:-8])
				l.recv_mac(r[-8:])

				assert pkt.startswith(b'reqreset')

				await self.put(l.send_enc(os.urandom(16)) +
				    l.send_mac(8))

				r = await self.get()
				c = l.recv_enc(r[:-8])
				l.recv_mac(r[-8:])

				assert c == b'confirm'

				await self.put(l.send_enc(b'confirmed') +
				    l.send_mac(8))

				r = await self.get()
				cmd = l.recv_enc(r[:-8])
				l.recv_mac(r[-8:])

				assert cmd[0] == CMD_WAITFOR
				assert int.from_bytes(cmd[1:], byteorder='little') == 30

				await self.put(l.send_enc(cmd[0:1]) +
				    l.send_mac(8))

				r = await self.get()
				cmd = l.recv_enc(r[:-8])
				l.recv_mac(r[-8:])

				assert cmd[0] == CMD_RUNFOR
				assert int.from_bytes(cmd[1:5], byteorder='little') == 1
				assert int.from_bytes(cmd[5:], byteorder='little') == 50

				await self.put(l.send_enc(cmd[0:1]) +
				    l.send_mac(8))

				r = await self.get()
				cmd = l.recv_enc(r[:-8])
				l.recv_mac(r[-8:])

				assert cmd[0] == CMD_TERMINATE

				await self.put(l.send_enc(cmd[0:1]) +
				    l.send_mac(8))

		tsd = TestSD()
		l = LORANode(tsd)

		await l.start()

		await l.waitfor(30)

		await l.runfor(1, 50)

		await l.terminate()

		await tsd.drain()

		# Make sure all messages have been processed
		self.assertTrue(tsd.sendq.empty())
		self.assertTrue(tsd.recvq.empty())

		print('done')