fl0w-old/Wallaby/Wallaby.py

from Highway import Route, Pipe, Client
from Meh import Config, Option, ExceptionInConfigError
import Logging
import Utils

import socket
import time
import os
import sys
import subprocess
from random import randint
from _thread import start_new_thread
from ctypes import cdll
import threading
import json
import re

from behem0th import Client as SyncClient
from behem0th import EventHandler

CHANNEL = 2
IS_WALLABY = Utils.is_wallaby()

USERS_LOCATION = "/home/root/Documents/KISS/users.json"
LIB_WALLABY = "/usr/lib/libwallaby.so"


if not IS_WALLABY:
	Logging.warning("Binaries that were created for Wallaby Controllers will not run on a simulated Wallaby.")


def get_users():
	if IS_WALLABY:
		return list(json.loads(open(USERS_LOCATION, "r").read()).keys())
	return ["Default User"]


class SensorReadout:
	ANALOG = 1
	DIGITAL = 2
	NAMED_MODES = {ANALOG : "analog", DIGITAL : "digital"}
	MODES = tuple(NAMED_MODES.keys())


	def __init__(self, handler, poll_rate=0.2):
		self.poll_rate = poll_rate
		self.handler = handler
		self.peer_lock = threading.Lock()
		self.peers = {}
		self.readout_required = {SensorReadout.ANALOG : [],
			SensorReadout.DIGITAL : []}
		self.generate_random_values = False
		# Running on actual hardware?
		if IS_WALLABY:
			if not os.path.exists(LIB_WALLABY):
				Logging.error("The Wallaby library should normally exist on "
					"a Wallaby. You broke something, mate.")
				self.generate_random_values = True
		else:
			Logging.warning("Sensor data can not be read on a dev-system. "
				"Generating random values instead.")
			self.generate_random_values = True
		# Generate random values?
		if not self.generate_random_values:
			self.wallaby_library = cdll.LoadLibrary(LIB_WALLABY)
			self.get_sensor_value = self.__get_sensor_value
		else:
			self.get_sensor_value = self.__get_random_value
		start_new_thread(self._sensor_fetcher, ())


	def subscribe(self, port, mode, peer):
		if port not in self.readout_required[mode]:
			self.readout_required[mode].append(port)
		if not peer in self.peers:
			self.peer_lock.acquire()
			self.peers[peer] = {SensorReadout.ANALOG : [],
			SensorReadout.DIGITAL : []}
			self.peer_lock.release()
		self.peers[peer][mode].append(port)


	def unsubscribe(self, port, mode, peer):
		if peer in self.peers:
			if port in self.peers[peer][mode]:
				self.peer_lock.acquire()
				del self.peers[peer][mode][self.peers[peer][mode].index(port)]
				self.peer_lock.release()
				self.determine_required_readouts()


	def determine_required_readouts(self):
		readout_required = {SensorReadout.ANALOG : [],
			SensorReadout.DIGITAL : []}
		self.peer_lock.acquire()
		for peer in self.peers:
			for mode in (SensorReadout.ANALOG, SensorReadout.DIGITAL):
				for port in self.peers[peer][mode]:
					if not port in readout_required[mode]:
						readout_required[mode].append(port)
		self.peer_lock.release()
		self.readout_required = readout_required


	def unsubscribe_all(self, peer):
		if peer in self.peers:
			self.peer_lock.acquire()
			del self.peers[peer]
			self.peer_lock.release()
			self.determine_required_readouts()


	def __get_sensor_value(self, port, mode):
		if mode == SensorReadout.ANALOG:
			return self.wallaby_library.analog(port)
		elif mode == SensorReadout.DIGITAL:
			return self.wallaby_library.digital(port)


	def __get_random_value(self, port, mode):
		if mode == SensorReadout.ANALOG:
			return randint(0, 4095)
		elif mode == SensorReadout.DIGITAL:
			return randint(0, 1)



	def _sensor_fetcher(self):
		while True:
			current_values = {SensorReadout.ANALOG : {},
			SensorReadout.DIGITAL : {}}
			for mode in SensorReadout.MODES:
				for port in self.readout_required[mode]:
					current_values[mode][port] = self.get_sensor_value(port, mode)
			self.peer_lock.acquire()
			for peer in self.peers:
				readouts = 0
				response = {"analog" : {}, "digital" : {}}
				for mode in SensorReadout.NAMED_MODES:
					for port in self.peers[peer][mode]:
						response[SensorReadout.NAMED_MODES[mode]][port] = current_values[mode][port]
						readouts += 1
				if readouts != 0:
					self.handler.pipe(response, "sensor", peer)
			self.peer_lock.release()
			time.sleep(self.poll_rate)

	@staticmethod
	def valid_port(port, mode):
		if mode == SensorReadout.ANALOG:
			return port >= 0 and port <= 5
		elif mode == SensorReadout.DIGITAL:
			return port >= 0 and port <= 9
		return False


class Identify(Pipe):
	def run(self, data, peer, handler):
		Utils.play_sound(config.identify_sound)
		if handler.debug:
			Logging.success("I was identified!")



class ListPrograms(Pipe):
	def run(self, data, peer, handler):
		programs = []
		if os.path.isdir(PATH):
			for program in os.listdir(PATH):
				if "botball_user_program" in os.listdir(PATH + program):
					programs.append(program)
		else:
			Logging.error("Harrogate folder structure does not exist. "
				"You broke something, mate.")
		handler.pipe(programs, handler.reverse_routes[self], peer)


class StopPrograms(Pipe):
	NO_PROGRAMS_RUNNING = 1

	def run(self, data, peer, handler):
		if handler.debug:
			Logging.info("Stopping all botball programs.")
		if subprocess.call(["killall", "botball_user_program"],
			stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT):
			handler.pipe(self.__class__.NO_PROGRAMS_RUNNING, "stop_programs", peer)


class RunProgram(Pipe):
	PROGRAM_NOT_FOUND = 1

	def __init__(self, output_unbuffer):
		self.command = [output_unbuffer, "-i0", "-o0", "-e0"]


	def run(self, data, peer, handler):
		if type(data) is str:
			if data[-1] != "/":
				data = data + "/"
			path = "%s%s/botball_user_program" % (PATH, data)
			if os.path.isfile(path):
				program = subprocess.Popen(self.command + [path],
					stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
				start_new_thread(self.stream_stdout, (program, peer, handler))
			else:
				handler.pipe(self.__class__.PROGRAM_NOT_FOUND, "run_program", peer)
				if handler.debug:
					Logging.warning("Program '%s' not found." % data[:-1])

	def stream_stdout(self, program, peer, handler):
		# Poll process for new output until finished
		for line in iter(program.stdout.readline, b""):
			handler.pipe(line.decode(), "std_stream", peer)
		program.wait()
		exit_code = program.returncode if type(program.returncode) is int else -1
		handler.pipe({"exit_code" : exit_code}, "std_stream", peer)


class Sensor(Pipe):
	"""
	->
	{"subscribe" : {"analog" : [1, 2, 3], "digital" : [1, 2, 3]}}
	{"unsubscribe" : {"analog" : [1, 2, 3], "digital" : [1, 2, 3]}}
	{"poll_rate" : 10}

	<-
	{"analog" : {1 : 1240}, "digital" : {1 : 0, 2 : 0}}
	"""
	def run(self, data, peer, handler):
		if type(data) is dict:
			if "poll_rate" in data:
				if type(data["poll_rate"]) in (int, float) and data["poll_rate"] >= 0.1:
					self.sensor_readout.poll_rate = data["poll_rate"]
			for event in ("subscribe", "unsubscribe"):
				if event in data:
					for mode in ("analog", "digital"):
						if mode in data[event]:
							for port in data[event][mode]:
								if type(port) is int:
									if mode == "analog":
										mode = SensorReadout.ANALOG
									elif mode == "digital":
										mode = SensorReadout.DIGITAL
									if SensorReadout.valid_port(port, mode):
										if event == "subscribe":
											self.sensor_readout.subscribe(port, mode, peer)
										elif event == "unsubscribe":
											self.sensor_readout.unsubscribe(port, mode, peer)
		elif type(data) is str:
			if data == "unsubscribe":
				self.sensor_readout.unsubscribe_all(peer)


	def start(self, handler):
		self.sensor_readout = SensorReadout(handler)


class BinaryReceived(Pipe, EventHandler):
	def on_modified(self, event):
		print(event)


	def on_created(self, event):
		print(event)


	def start(self, handler):
		self.handler = handler



class Shutdown(Pipe):
	def run(self, data, peer, handler):
		try:
			subprocess.check_output(["shutdown", "now"])
		except subprocess.CalledProcessError:
			Logging.warning("Shutdown could not be initiated.")


class Reboot(Pipe):
	def run(self, data, peer, handler):
		try:
			subprocess.check_output(["reboot"])
		except subprocess.CalledProcessError:
			Logging.warning("Reboot could not be initiated.")


class Subscribe(Route):
	def start(self, handler):
		handler.send({"name" : Utils.get_hostname(), "channel" : CHANNEL}, "subscribe")


class WhoAmI(Route):
	def run(self, data, handler):
		Logging.header("ID: '%s' - Running as '%s'" % (data["id"], data["user"]))

	def start(self, handler):
		handler.send(None, "whoami")


class Hostname(Pipe):
	def run(self, data, peer, handler):
		if type(data) is dict:
			if "set" in data:
				try:
					Utils.set_hostname(str(data["set"]))
				except Utils.HostnameNotChangedError:
					if IS_WALLABY:
						Logging.error("Hostname change unsuccessful. "
							"Something is wrong with your Wallaby.")
					else:
						Logging.warning("Hostname change unsuccessful. "
							"This seems to be a dev-system, so don't worry too "
							"much about it.")


class Processes(Pipe):
	def run(self, data, peer, handler):
		handler.pipe(
			subprocess.check_output(["ps", "aux"]).decode().split("\n")[1:-1],
			"processes", peer)


class Output(Pipe):
	def __init__(self):
		self.content = []
		self.handler = None
		self.subscribers = []

		self.escape_regex = re.compile(r"(\[\d+m?)")

		self.stdout = Output.Std(Logging.stdout, self.push)
		self.stderr = Output.Std(Logging.stderr, self.push)

		Logging.stdout = self.stdout
		Logging.stderr = self.stderr

	def run(self, data, peer, handler):
		if data == "subscribe":
			self.subscribers.append(peer)
		elif data == "unsubscribe":
			self.unsubscribe(peer)


	def start(self, handler):
		self.handler = handler


	def push(self, s):
		s = self.escape_regex.sub("", s)
		for peer in self.subscribers:
			self.handler.pipe(s, "output", peer)


	def unsubscribe(self, peer):
		if peer in self.subscribers:
			del self.subscribers[self.subscribers.index(peer)]		

	class Std:
		def __init__(self, old_std, write_callback):
			self.old_std = old_std
			self.write_callback = write_callback


		def write(self, s):
			self.old_std.write(s)
			self.write_callback(s)


		def flush(self):
			self.old_std.flush()


class Handler(Client):
	def setup(self, routes, debug=False):
		super().setup(routes, debug_suppress_routes=["pipe"], debug=debug)


	def peer_unavaliable(self, peer):
		if self.debug:
			Logging.info("Unsubscribing '%s' from all sensor updates." % peer)
		self.routes["sensor"].sensor_readout.unsubscribe_all(peer)
		self.routes["output"].unsubscribe(peer)


if IS_WALLABY:
	if not "fl0w" in get_users():
		json.loads(open(USERS_LOCATION, "r").read())["fl0w"] = {"mode" : "Advanced"}
		try:
			os.mkdir(FL0W_USER_PATH)
		except FileExistsError:
			pass
	PATH = FL0W_USER_PATH + "bin/"
else:
	if len(sys.argv) == 2:
		if os.path.exists(sys.argv[1]):
			PATH = os.path.abspath(sys.argv[1])
		else:
			Logging.error("Location does not exist.")
			exit(1)
	else:
		Logging.error("Location has to be provided in dev-env.")
		exit(1)

if PATH[-1] != "/":
	PATH = PATH + "/"



CONFIG_PATH = "wallaby.cfg"

config = Config()
config.add(Option("server_address", "ws://127.0.0.1:3077"))
config.add(Option("debug", True, validator=lambda x: True if True or False else False))
config.add(Option("output_unbuffer", "stdbuf"))
config.add(Option("identify_sound", "Wallaby/identify.wav",
	validator=lambda sound: os.path.isfile(sound)))


try:
    config = config.load(CONFIG_PATH)
except (IOError, ExceptionInConfigError):
    config.dump(CONFIG_PATH)
    config = config.load(CONFIG_PATH)


try:
	ws = Handler(config.server_address)
	# setup has to be called before the connection is established
	ws.setup({"subscribe" : Subscribe(), "hostname" : Hostname(),
		"processes" : Processes(), "sensor" : Sensor(),
		"identify" : Identify(), "list_programs" : ListPrograms(),
		"whoami" : WhoAmI(), "run_program" : RunProgram(config.output_unbuffer),
		"stop_programs" : StopPrograms(), "shutdown" : Shutdown(),
		"reboot" : Reboot(), "binary_received" : BinaryReceived(),
		"output" : Output()},
		debug=config.debug)
	ws.connect()
	sync_client = SyncClient(path=PATH, verbose_log=config.debug,
		event_handler=ws.routes["binary_received"])
	sync_client.connect(Utils.get_ip_from_url(config.server_address))
	ws.run_forever()
except KeyboardInterrupt:
	ws.close()
	#sync_client.close()