From 9a68cdad2a0d359878fd3ab498d9396660d9da0c Mon Sep 17 00:00:00 2001
From: csucu <csucu@github.com>
Date: Mon, 16 Jan 2017 12:22:19 +0000
Subject: [PATCH] Automatically reconnect to the Tor control port

---
 onionbalance/instance.py | 23 ++++++++++++++++++++---
 onionbalance/service.py  | 20 +++++++++++++-------
 onionbalance/util.py     | 16 ++++++++++++++++
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/onionbalance/instance.py b/onionbalance/instance.py
index 9bdcadf..bb4cd1f 100644
--- a/onionbalance/instance.py
+++ b/onionbalance/instance.py
@@ -6,6 +6,7 @@
 
 from onionbalance import log
 from onionbalance import config
+from onionbalance import util
 
 logger = log.get_logger()
 
@@ -18,12 +19,25 @@ def fetch_instance_descriptors(controller):
 
     # Clear Tor descriptor cache before making fetches by sending NEWNYM
     # pylint: disable=no-member
-    controller.signal(stem.control.Signal.NEWNYM)
-    time.sleep(5)
 
+    while True:
+        try:
+            controller.signal(stem.control.Signal.NEWNYM)
+            time.sleep(5)
+            break
+        except stem.SocketClosed:
+            logger.error("Failed to send NEWNYM signal, socket is closed.")
+            util.reauthenticate(controller, logger)
+   
     for service in config.services:
         for instance in service.instances:
-            instance.fetch_descriptor()
+            while True:
+                try:
+                    instance.fetch_descriptor()
+                    break
+                except stem.SocketClosed:
+                    logger.error("Failed to fecth descriptor, socket is closed")
+                    util.reauthenticate(controller, logger)
 
 
 class Instance(object):
@@ -66,6 +80,9 @@ def fetch_descriptor(self):
         try:
             self.controller.get_hidden_service_descriptor(self.onion_address,
                                                           await_result=False)
+        except stem.SocketClosed:
+            # Tor maybe restarting.
+            raise
         except stem.DescriptorUnavailable:
             # Could not find the descriptor on the HSDir
             self.received = None
diff --git a/onionbalance/service.py b/onionbalance/service.py
index a96d224..bc2823a 100644
--- a/onionbalance/service.py
+++ b/onionbalance/service.py
@@ -219,18 +219,24 @@ def _publish_descriptor(self, deviation=0):
     def _upload_descriptor(self, signed_descriptor, replica, hsdirs=None):
         """
         Convenience method to upload a descriptor
-
         Handle some error checking and logging inside the Service class
         """
         if hsdirs and not isinstance(hsdirs, list):
             hsdirs = [hsdirs]
 
-        try:
-            descriptor.upload_descriptor(self.controller, signed_descriptor,
-                                         hsdirs=hsdirs)
-        except stem.ControllerError:
-            logger.exception("Error uploading descriptor for service "
-                             "%s.onion.", self.onion_address)
+        while True:
+            try:
+                descriptor.upload_descriptor(self.controller, signed_descriptor,
+                                             hsdirs=hsdirs)
+                break
+            except stem.SocketClosed:
+                logger.error("Error uploading descriptor for service "
+                                 "%s.onion, Socket is closed.", self.onion_address)
+                util.reauthenticate(self.controller, logger)
+            except stem.ControllerError:
+                logger.exception("Error uploading descriptor for service "
+                                 "%s.onion.", self.onion_address)
+                break
 
     def descriptor_publish(self, force_publish=False):
         """
diff --git a/onionbalance/util.py b/onionbalance/util.py
index 3b91192..e606a80 100644
--- a/onionbalance/util.py
+++ b/onionbalance/util.py
@@ -6,10 +6,14 @@
 import base64
 import binascii
 import os
+import stem
+import time
 
 # import Crypto.Util
 import Crypto.PublicKey
 
+from onionbalance import config
+
 
 def add_pkcs1_padding(message):
     """Add PKCS#1 padding to **message**."""
@@ -160,3 +164,15 @@ def is_directory_empty(path):
         return False
     else:
         return True
+
+
+def reauthenticate(controller,logger):
+    """
+    Tries to authenticate to the controller 
+    """
+    time.sleep(10)
+    try: 
+        controller.authenticate(password=config.TOR_CONTROL_PASSWORD)
+    except stem.connection.AuthenticationFailure:
+        logger.error("Failed to re-authenticate controller.")
+        
\ No newline at end of file