From 85212fb1710ff7811312401ddba4af0c90f1f4db Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 23 Feb 2019 17:49:40 -0700
Subject: [PATCH] Add retry option when connecting to osTicket

* Fixes issue #51
---
 .bin/Scripts/functions/osticket.py | 39 ++++++++++++++++++------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/.bin/Scripts/functions/osticket.py b/.bin/Scripts/functions/osticket.py
index 0c051974..ca7cf153 100644
--- a/.bin/Scripts/functions/osticket.py
+++ b/.bin/Scripts/functions/osticket.py
@@ -18,6 +18,11 @@ REGEX_TEMPS = re.compile(r'^\s*(.*?)\s+(idle:)(.*)$')
 REGEX_SENSOR = re.compile(r'^(.*?)(\s*)$')
 
 
+# Error Classes
+class osTicketConnectionError(Exception):
+  pass
+
+
 # Classes
 class osTicket():
   """Class to track osTicket data and functions."""
@@ -70,15 +75,9 @@ class osTicket():
         # Connection established
         break
 
-    # Disable if necessary
-    if self.db_cursor is None:
-      self.disabled = True
-      self.tunnel_proc.kill()
-      if not silent:
-        print_warning('Failed to connect to osTicket')
-        print_standard('Integration disabled for this session')
-        print_standard(' ')
-        pause()
+    # Raise exception unless silenced
+    if self.db_cursor is None and not silent:
+      raise osTicketConnectionError('Failed to connect to osTicket.')
 
   def convert_report(self, name, test):
     """Convert report into an osTicket friendly format, returns list."""
@@ -347,7 +346,7 @@ class osTicket():
   def get_flag(self, ticket_id, flag_name):
     """Get flag in osTicket."""
     flag_value = None
-    self.connect(silent=True)
+    self.connect()
 
     # Bail if disabled
     if self.disabled:
@@ -395,11 +394,21 @@ class osTicket():
   def get_ticket_number(self):
     """Get ticket number and confirm with name from osTicket DB."""
     ticket_number = None
-    self.connect(silent=False)
 
-    # Bail if disabled
-    if self.disabled:
-      return None
+    # Connect
+    while True:
+      try:
+        self.connect(silent=False)
+      except osTicketConnectionError:
+        print_warning('Failed to connect to osTicket')
+        if not ask('Try again?'):
+          print_standard('Integration disabled for this session')
+          self.disabled = True
+          self.tunnel_proc.kill()
+          return None
+      else:
+        # Connection successful
+        break
 
     # Main loop
     while ticket_number is None:
@@ -498,7 +507,7 @@ class osTicket():
 
   def set_flag(self, ticket_id, flag_name, flag_value):
     """Set flag in osTicket."""
-    self.connect(silent=True)
+    self.connect()
 
     # Bail if disabled
     if self.disabled: