Include disconnect and send function on client side

To disconnect the socket or send a message we use in different parts of the code
the attribute socket which is declared inside of the client. This way will
create us important problem in the future if we want to change the behavior in
our client sockets. This patch adds two new methods inside the client and we
can use them in other classes giving us the easy way to send messages and
disconnect the client.

2 files changed, 26 insertions, 16 deletions

diff --git a/src/ogClient.py b/src/ogClient.py
index c1a911d..48ae5e7 100644
--- a/src/ogClient.py
+++ b/src/ogClient.py
@@ -44,6 +44,10 @@ class ogClient:
 

 			print ('Error connect ' + str(err))

 

+	def send(self, msg):

+		self.sock.send(bytes(msg, 'utf-8'))

+		return len(msg)

+

 	def connect2(self):

 		try:

 			self.sock.connect((self.ip, self.port))

@@ -85,14 +89,17 @@ class ogClient:
 

 		if self.trailer and len(self.data) >= self.content_len:

 			httpparser.parser(self.data)

-			if not ogprocess.processOperation(httpparser.getRequestOP(), httpparser.getURI(), self.sock):

-				self.state = State.FORCE_DISCONNECTED

+			ogprocess.processOperation(httpparser.getRequestOP(), httpparser.getURI(), self)

 

 			# Cleanup state information from request

 			self.data = ""

 			self.content_len = 0

 			self.trailer = False

 

+	def disconnect(self):

+		self.state = State.FORCE_DISCONNECTED

+		self.sock.close()

+

 	def run(self):

 		while 1:

 			sock = self.get_socket()

diff --git a/src/ogProcess.py b/src/ogProcess.py
index def76f1..27954a7 100644
--- a/src/ogProcess.py
+++ b/src/ogProcess.py
@@ -7,33 +7,36 @@ if platform.system() == 'Linux':
 	from src.linux import ogOperations
 
 class ogProcess():
-	def processOperation(self, op, URI, sock):
+	def processOperation(self, op, URI, client):
 		if ("poweroff" in URI):
-			sock.send(bytes(ogRest.getResponse(ogRest.ogResponses.IN_PROGRESS), 'utf-8'))
-			sock.close()
-			self.process_poweroff()
-			return 0
+			self.process_poweroff(client)
 		elif ("reboot" in URI):
-			sock.send(bytes(ogRest.getResponse(ogRest.ogResponses.IN_PROGRESS), 'utf-8'))
-			sock.close()
-			self.process_reboot()
-			return 0
+			self.process_reboot(client)
 		elif ("probe" in URI):
-			sock.send(bytes(ogRest.getResponse(ogRest.ogResponses.OK), 'utf-8'))
+			self.process_probe(client)
 		else:
-			sock.send(bytes(ogRest.getResponse(ogRest.ogResponses.BAD_REQUEST), 'utf-8'))
+			client.send(ogRest.getResponse(ogRest.ogResponses.BAD_REQUEST))
 
-		return 1
+		return 0
 
-	def process_reboot(self):
+	def process_reboot(self, client):
 		# Rebooting thread
 		def rebt():
 			ogOperations.reboot()
+
+		client.send(ogRest.getResponse(ogRest.ogResponses.IN_PROGRESS))
+		client.disconnect()
 		threading.Thread(target=rebt).start()
 
-	def process_poweroff(self):
+	def process_poweroff(self, client):
 		# Powering off thread
 		def pwoff():
 			time.sleep(2)
 			ogOperations.poweroff()
+
+		client.send(ogRest.getResponse(ogRest.ogResponses.IN_PROGRESS))
+		client.disconnect()
 		threading.Thread(target=pwoff).start()
+
+	def process_probe(self, client):
+		client.send(ogRest.getResponse(ogRest.ogResponses.OK))