Automatically reconnect to the Tor control port.

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,26 @@ 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 +81,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

onionbalance/service.py

@@ -219,18 +219,25 @@ 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):
         """

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,14 @@ 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.")