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()