This repository has been archived on 2025-06-04. You can view files and clone it, but you cannot make any changes to it's state, such as pushing and creating new issues, pull requests or comments.
fl0w-old/Wallaby/Wallaby.py
2016-12-05 22:55:50 +01:00

244 lines
No EOL
7.2 KiB
Python

from Highway import Route, Pipe, Client
import Logging
import Config
from Utils import is_wallaby, set_hostname, get_hostname
import socket
import time
import os
import sys
import subprocess
import _thread
CHANNEL = 2
IS_WALLABY = is_wallaby()
PATH = "/home/root/Documents/KISS/bin/" if IS_WALLABY else (sys.argv[1] if len(sys.argv) > 1 else None)
LIB_WALLABY = "/usr/lib/libwallaby.so"
if not PATH:
Logging.error("No path specified. (Necessary on simulated Wallaby controllers.)")
exit(1)
if not IS_WALLABY:
Logging.warning("Binaries that were created for Wallaby Controllers will not run on a simulated Wallaby.")
class SensorReadout:
ANALOG = 1
DIGITAL = 2
def __init__(self, poll_rate=0.2, handler):
self.poll_rate = poll_rate
self.handler = handler
self.peers = {}
self.generate_random_values = False
if Utils.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
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, peer, mode):
if not peer in self.peers:
self.peers[peer] = {SensorReadout.ANALOG : [], SensorReadout.DIGITAL : []}
else:
self.peers[peer][mode].append(port)
def unsubscribe(self, port, peer, mode):
if peer in self.peers:
if port in self.peers[peer][mode]:
del self.peers[peer][mode][self.peers[peer][mode].index(port)]
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 == DIGITAL:
return randint(0, 1)
def _sensor_fetcher(self, handler):
while True:
current_values = {"analog" : {}, "digital" : {}}
for peer in self.peers:
response = {"analog" : {}, "digital" : {}}
for port in self.peers[peer][SensorReadout.ANALOG]:
if not port in current_values["analog"]:
sensor_value = get_sensor_value(port, ANALOG)
response["analog"][port] = sensor_value
current_values["analog"][port] = sensor_value
else:
response["analog"][port] = current_values["analog"][port]
for port in self.peers[peer][SensorReadout.DIGITAL]:
if not port in current_values["digital"]:
sensor_value = get_sensor_value(port, DIGITAL)
response["digital"][port] = sensor_value
current_values["digital"][port] = sensor_value
else:
response["digital"][port] = current_values["digital"][port]
self.handler.pipe(response, "sensor", peer)
time.sleep(self.poll_rate)
class Sensor(Pipe):
"""
{"subscribe" : {"analog" : [1, 2, 3], "digital" : [1, 2, 3]}}
{"unsubscribe" : {"analog" : [1, 2, 3], "digital" : [1, 2, 3]}}
"""
def run(self, data, peer, handler):
if type(data) is dict:
if "subscribe" in data:
if "analog" in data["subscribe"]:
for port in data["subscribe"]["analog"]:
if type(port) is int:
if port >= 0 and port <= 10:
sensor_readout.subscribe(port, handler,
SensorReadout.ANALOG)
if "digital" in data["subscribe"]:
for port in data["subscribe"]["digital"]:
if type(port) is int:
if port >= 0 and port <= 10:
sensor_readout.subscribe(port, handler,
SensorReadout.DIGITAL)
if "unsubscribe" in data:
if "analog" in data["unsubscribe"]:
for port in data["unsubscribe"]["analog"]:
if type(port) is int:
if port >= 0 and port <= 10:
sensor_readout.unsubscribe(port, handler,
SensorReadout.ANALOG)
if "digital" in data["unsubscribe"]:
for port in data["unsubscribe"]["digital"]:
if type(port) is int:
if port >= 0 and port <= 10:
sensor_readout.unsubscribe(port, handler,
SensorReadout.DIGITAL)
class WallabyControl(Route):
def __init__(self, output_unbuffer):
self.output_unbuffer = output_unbuffer
self.actions_with_params = {"run" : self.run_program}
self.actions_without_params = {"disconnect" : self.disconnect,
"reboot" : self.reboot, "shutdown" : self.shutdown, "stop" : self.stop}
self.currently_running_program = None
def run(self, data, handler):
if type(data) is str:
if data in self.actions_without_params.keys():
self.actions_without_params[data](handler)
elif type(data) is dict:
for action in data:
if action in self.actions_with_params.keys():
_thread.start_new_thread(self.actions_with_params[action], (handler, data[action]))
def run_program(self, handler, program):
command = [self.output_unbuffer, "-i0", "-o0", "-e0"]
command.append("%s%s/botball_user_program" % (handler.sync.folder, program))
self.currently_running_program = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
# Poll process for new output until finished
for line in iter(self.currently_running_program.stdout.readline, b""):
handler.sock.send(line.decode(), "std_stream")
self.currently_running_program.wait()
handler.sock.send({"return_code" : self.currently_running_program.returncode}, "std_stream")
self.currently_running_program = None
def stop(self, handler):
if self.currently_running_program != None:
Logging.info("Killing currently running programm.")
self.currently_running_program.kill()
else:
Logging.info("No program started by fl0w.")
def reboot(self, handler):
self.disconnect(handler)
os.system("reboot")
exit(0)
def shutdown(self, handler):
self.disconnect(handler)
os.system("shutdown -h 0")
def disconnect(self, handler):
self.stop(handler)
handler.sock.close()
class Subscribe(Route):
def start(self, handler):
handler.send({"name" : get_hostname(), "channel" : CHANNEL}, "subscribe")
class Hostname(Pipe):
def run(self, data, peer, handler):
if type(data) is dict:
if "set" in data:
set_hostname(str(data["set"]))
class Processes(Pipe):
def run(self, data, peer, handler):
handler.pipe(
subprocess.check_output(["ps", "aux"]).decode().split("\n")[1:-1],
"processes", peer)
class Handler(Client):
def setup(self, routes, debug=False):
super().setup(routes, piping=True, debug=debug)
CONFIG_PATH = "wallaby.cfg"
config = Config.Config()
config.add(Config.Option("server_address", "ws://127.0.0.1:3077"))
config.add(Config.Option("debug", False, validator=lambda x: True if True or False else False))
config.add(Config.Option("output_unbuffer", "stdbuf"))
try:
config = config.read_from_file(CONFIG_PATH)
except FileNotFoundError:
config.write_to_file(CONFIG_PATH)
config = config.read_from_file(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()},
debug=config.debug)
ws.connect()
ws.run_forever()
except KeyboardInterrupt:
ws.close()