From 5b5827f26e27761781039bcfad5f45e05c048bcf Mon Sep 17 00:00:00 2001
From: Jesse Luehrs <doy@tozt.net>
Date: Fri, 17 Oct 2014 16:42:53 -0400
Subject: reconnect to the server when the connection drops

---
 termcast_client/__init__.py | 44 ++++++++++++++++++++++++++++++++++++++------
 1 file changed, 38 insertions(+), 6 deletions(-)

(limited to 'termcast_client/__init__.py')

diff --git a/termcast_client/__init__.py b/termcast_client/__init__.py
index 70d1798..b3c2b4b 100644
--- a/termcast_client/__init__.py
+++ b/termcast_client/__init__.py
@@ -7,6 +7,7 @@ import signal
 import socket
 import ssl
 import sys
+import time
 
 from . import pity
 from . import py2compat
@@ -21,11 +22,7 @@ class Client(object):
         self.fingerprint = fingerprint
 
     def run(self, argv):
-        self.sock = socket.socket()
-        self.sock.connect((self.host, self.port))
-        if self.tls:
-            self._starttls()
-        self.sock.send(self._build_connection_string())
+        self._new_socket()
         self.winch_set = False
         pity.spawn(
             argv,
@@ -33,6 +30,30 @@ class Client(object):
             handle_window_size=True
         )
 
+    def _new_socket(self):
+        self.sock = socket.socket()
+        self.sock.connect((self.host, self.port))
+        if self.tls:
+            self._starttls()
+        self.sock.send(self._build_connection_string())
+
+    def _renew_socket(self):
+        self.sock.close()
+
+        while True:
+            print("Disconnected from server, reconnecting...")
+            time.sleep(5)
+            try:
+                self.sock = socket.socket()
+                self.sock.connect((self.host, self.port))
+                if self.tls:
+                    self._starttls()
+                self.sock.send(self._build_connection_string())
+                return
+            except:
+                pass
+
+
     def _starttls(self):
         self.sock.send(b'starttls\n')
         context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
@@ -52,7 +73,18 @@ class Client(object):
             self.winch_set = True
 
         data = os.read(fd, 1024)
-        self.sock.send(data)
+        to_send = data
+        datalen = len(to_send)
+        while datalen > 0:
+            try:
+                sent = self.sock.send(to_send)
+                if sent == 0:
+                    raise "disconnect"
+                to_send = to_send[sent:]
+                datalen = len(to_send)
+            except:
+                self._renew_socket()
+
         return data
 
     def _winch(self, signum, frame):
-- 
cgit v1.2.3-54-g00ecf