From 6bc587c5757800d3cc0509b46ba45cd1a9dd7c4c Mon Sep 17 00:00:00 2001
From: Philip Trauner <philip.trauner@aol.com>
Date: Sat, 20 Feb 2016 10:04:09 +0100
Subject: [PATCH] Added reboot support to Wallaby Control

---
 Server/Server.py     |  6 ++++--
 Sublime/fl0w/fl0w.py | 43 ++++++++++---------------------------------
 Wallaby/Wallaby.py   |  6 +++++-
 3 files changed, 19 insertions(+), 36 deletions(-)

diff --git a/Server/Server.py b/Server/Server.py
index ecd6427..f575cb1 100644
--- a/Server/Server.py
+++ b/Server/Server.py
@@ -27,11 +27,13 @@ class WallabyControl(Routing.ServerRoute):
 			for wallaby in handler.broadcast.channels[Handler.Channels.WALLABY]:
 				address_pair = "%s:%d" % (wallaby.address, wallaby.port)
 				if address_pair in data.keys():
-					if data[address_pair] in ("stop", "restart", "disconnect"):
+					if data[address_pair] in ("stop", "restart", "disconnect", "reboot"):
 						wallaby.send(data[address_pair], "wallaby_control")
 					elif type(data[address_pair]) is dict:
 						if "run" in data[address_pair]:
 							Logging.warning("Remote binary execution not yet implemented. (file_sync and compile required)")
+					else:
+						Logging.warning("'%s:%d' has issued an invalid control command." % (handler.info[0], handler.info[1]))
 				return
 			handler.sock.send("Wallaby not connected anymore.", "error_report")
 
@@ -78,7 +80,7 @@ class Handler(Server.Handler):
 		Logging.info("'%s:%d' disconnected." % (self.sock.address, self.sock.port))
 
 
-server = Server(("127.0.0.1", 3077), debug=True)
+server = Server(("172.20.10.3", 3077), debug=True)
 
 broadcast = Broadcast()
 # Populating broadcast channels with all channels defined in Handler.Channels
diff --git a/Sublime/fl0w/fl0w.py b/Sublime/fl0w/fl0w.py
index 86c5385..4d7b76c 100644
--- a/Sublime/fl0w/fl0w.py
+++ b/Sublime/fl0w/fl0w.py
@@ -26,46 +26,23 @@ import webbrowser
 def plugin_unloaded():
 	observer.stop()
 	try:
-		fl0w.sock.close()
-		fl0w.connected = False
+		fl0w.invoke_disconnect()
 	except:
 		pass
 	print("Observer stopped!")
 
 
-class HandledESock:
-	def __init__(self, sock, disconnect):
-		self._sock = sock
-		self.disconnect = disconnect
-
-	def __getattr__(self, attr):
-		if attr == "send":
-			return self.send
-		return getattr(self._sock, attr)
-
-	def send(self, data, route):
-		try:
-			self._sock.send(data, route)
-		except OSError:
-			self._sock.close()
-			self.disconnect()
-			_thread.exit()
-
-	def recv(self):
-		try:
-			return self._sock.recv()
-		except OSError:
-			self._sock.close()
-			self.disconnect()
-			_thread.exit()
-
 
 def sock_handler(sock, routes, handler):
 	sock.send("st", "set_type")
 	while 1:
-		data = sock.recv()
-		if data[1] in routes:
-			routes[data[1]].run(data[0], handler)
+		try:
+			data = sock.recv()
+			if data[1] in routes:
+				routes[data[1]].run(data[0], handler)
+		except (OSError, socket.error):
+			handler.invoke_disconnect()
+			break
 
 
 class ReloadHandler(FileSystemEventHandler):
@@ -127,7 +104,7 @@ class Fl0w:
 
 	def wallaby_control_submenu(self, wallaby):
 		menu = Menu(subtitles=False)
-		for command in ("Stop", "Restart", "Disconnect"):
+		for command in ("Stop", "Restart", "Reboot", "Disconnect"):
 			menu.add(Entry(command, action=self.sock.send, kwargs={"data" : {wallaby : command.lower()}, "route" : "wallaby_control"}))
 		menu.invoke(self.window)
 
@@ -167,7 +144,7 @@ class Fl0w:
 		connect_details = connect_details.split(":")
 		if len(connect_details) == 2:
 			try:
-				self.sock = HandledESock(ESock(socket.create_connection((connect_details[0], int(connect_details[1])))), self.invoke_disconnect)
+				self.sock = ESock(socket.create_connection((connect_details[0], int(connect_details[1]))), disconnect_callback=self.invoke_disconnect)
 				sublime.status_message("Connected to %s:%s." % (connect_details[0], connect_details[1]))
 				self.connected = True
 				_thread.start_new_thread(sock_handler, (self.sock, {"error_report": Fl0w.ErrorReport(), "info" : Fl0w.Info(), "wallaby_control" : Fl0w.WallabyControl()}, self))
diff --git a/Wallaby/Wallaby.py b/Wallaby/Wallaby.py
index c153c3f..b201d40 100644
--- a/Wallaby/Wallaby.py
+++ b/Wallaby/Wallaby.py
@@ -9,7 +9,8 @@ import sys
 
 class WallabyControl(Routing.ClientRoute):
 	def run(self, data, handler):
-		commands = {"stop" : self.stop, "restart" : self.restart, "disconnect" : self.disconnect}
+		commands = {"stop" : self.stop, "restart" : self.restart, 
+		"disconnect" : self.disconnect, "reboot" : self.reboot}
 		if data in commands:
 			commands[data](handler)
 
@@ -21,6 +22,9 @@ class WallabyControl(Routing.ClientRoute):
 		time.sleep(15)
 		os.execl(sys.executable, *([sys.executable]+sys.argv))
 
+	def reboot(self, handler):
+		os.system("reboot")
+
 	def disconnect(self, handler):
 		handler.sock.close()