fix(client): #69 batch command use async socket

client.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python
 
 import sys
+import time
 import socket
 import select
 import termios
@@ -10,36 +11,74 @@
 import os
 import struct
 
-def netfail(rhost, rport, bport, string):
-    sys.exit("host [%s], port [%d], bind port [%d] : %s" \
-                    % (rhost, rport, bport, string))
-
 
 class Client:
+    def _netfail(self, string):
+        self.close()
+        sys.exit("host [%s], port [%d], bind port [%d] : %s" \
+                        % (self.rhost, self.rport, self.bport, string))
+
     def __init__(self, rhost="127.0.0.1", rport=3000, bport=64832):
+	self.rhost = rhost
+	self.rport = rport
+	self.bport = bport
         try:
             self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+	    self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
             self.s.bind((rhost, bport))
             self.s.connect((rhost, rport))
         except socket.gaierror:
-            netfail(rhost, rport, bport, "Socket error")
+            self._netfail("Socket error")
         except OverflowError:
-            netfail(rhost, rport, bport, "Bind error")
+            self._netfail("Bind error")
         except:
-            netfail(rhost, rport, bport, "Connection refused")
+            self._netfail("Connection refused")
 
     @staticmethod
     def usage(name):
-        print "%s <address> <port> <bind_port> [password] [batch_file]" % name
+        print "%s <address> <port> <bind_port> [password] [batch_file | -c command]" % name
 
     def tty_raw_mode(self):
         self.old_attrs = termios.tcgetattr(sys.stdin)
         tty.setcbreak(sys.stdin)
 
+    def _recv_timeout(self, timeout=0.1):
+        # **NOTE** The socket need to be non blocking
+        total_data = [];
+        data = '';
+        begin = time.time()
+        while 1:
+            # if you got some data, then break after timeout
+            if total_data and time.time()-begin > timeout:
+                break
+            # if you got no data at all, wait a little longer, twice the timeout
+            elif time.time()-begin > timeout*2:
+                break
+            try:
+                data = self.s.recv(8192)
+                if data:
+                    total_data.append(data)
+                    begin = time.time()
+                else:
+                    time.sleep(0.1)
+            except:
+                pass
+        return ''.join(total_data)
+
+    def _print_response(self):
+        data = self._recv_timeout()
+        sys.stdout.write(data)
+        sys.stdout.flush()
+
     def password(self, passwd):
+        # Need to be run before run_command or loop_batch
+	self.s.setblocking(0)
         self.s.send(passwd + "\n")
-        data = self.s.recv(1024)
-        data = self.s.recv(1024)
+        self._recv_timeout()
+
+    def run_command(self, cmd):
+        self.s.send(cmd)
+        self._print_response()
 
     def connection(self, passwd=""):
         if passwd == "":
@@ -49,7 +88,11 @@ def connection(self, passwd=""):
         else:
             passwd += "\n"
         self.s.send(passwd)
-        data = self.s.recv(1024)
+        try:
+            data = self.s.recv(1024)
+        except (socket.error, socket.herror) as e:
+            (errno, error) = e
+	    self._netfail(error)
         if not data:
             sys.stdout.write("\n")
             self.close()
@@ -58,16 +101,9 @@ def connection(self, passwd=""):
         sys.stdout.flush()
 
     def loop_batch(self, batch, fd):
-        self.s.send("\n")
         for line in fd:
-            self.s.send(line)
-        data = self.s.recv(1024)
-        if data[0:2] == "\r\n":
-            data = data[2:]
-        sys.stdout.write(data)
-        while data:
-            data = self.s.recv(1024)
-            sys.stdout.write(data)
+            self.run_command(line)
+        sys.stdout.write("\n")
 
     def loop(self):
         while 1:
@@ -84,15 +120,19 @@ def loop(self):
                     sys.stdout.flush()
 
     def close(self):
-        self.s.close()
-        if hasattr(client, 'old_attrs'):
+        try:
+            self.s.shutdown(socket.SHUT_RDWR)
+            self.s.close()
+        except:
+            pass
+        if hasattr(self, 'old_attrs'):
             termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_attrs)
 
 
 if len(sys.argv) < 4:
     Client.usage(sys.argv[0])
     exit(1)
-elif len(sys.argv) >= 6:
+elif len(sys.argv) == 6:
     try:
         fd = open(sys.argv[5], "r")
     except IOError:
@@ -112,5 +152,7 @@ def close(self):
 elif len(sys.argv) == 6:
     client.password(sys.argv[4])
     client.loop_batch(sys.argv[5], fd)
+elif len(sys.argv) == 7:
+    client.password(sys.argv[4])
+    client.run_command(sys.argv[6] + "\n")
 client.close()
-

tests/quick/client/client.sh

@@ -5,7 +5,7 @@ cd `git rev-parse --show-toplevel`
 
 ! ./client.py
 ./client.py host port bind_port pass missing_file.txt | grep "cannot open"
-./client.py | grep "client.py <address> <port> <bind_port> \[password\] \[batch_file\]"
+./client.py | grep "client.py <address> <port> <bind_port> \[password\] \[batch_file | -c command\]"
 ./client.py beurk 3005 6666 2>&1 |
     grep "host \[beurk\], port \[3005\], bind port \[6666\] : Socket error"
 ./client.py 127.7.7.7 3005 77777 2>&1 |