Merge pull request #27 from spark1security/ms/issue26

Changed get_data() to automatically reconnect client

src/n0s1/controllers/asana_controller.py

@@ -10,11 +10,11 @@
 class AsanaController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
 
-    def set_config(self, config):
+    def set_config(self, config=None):
+        super().set_config(config)
         import asana
-        TOKEN = config.get("token", "")
+        TOKEN = self._config.get("token", "")
         self._client = asana.Client.access_token(TOKEN)
         return self.is_connected()
 
@@ -55,12 +55,15 @@ def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return {}
 
+        self.connect()
         if workspaces := self._client.workspaces.get_workspaces():
             for w in workspaces:
                 workspace_gid = w.get("gid", "")
+                self.connect()
                 if projects := self._client.projects.get_projects_for_workspace(workspace_gid):
                     for p in projects:
                         project_gid = p.get("gid", "")
+                        self.connect()
                         if tasks := self._client.tasks.get_tasks_for_project(project_gid, opt_fields=["name", "gid", "notes", "permalink_url"]):
                             for t in tasks:
                                 comments = []
@@ -69,6 +72,7 @@ def get_data(self, include_coments=False, limit=None):
                                 description = t.get("notes", "")
                                 url = t.get("permalink_url", "")
                                 if include_coments:
+                                    self.connect()
                                     if stories := self._client.stories.get_stories_for_task(task_gid):
                                         for s in stories:
                                             if s.get("type", "").lower() == "comment".lower():
@@ -80,6 +84,7 @@ def get_data(self, include_coments=False, limit=None):
     def post_comment(self, task_gid, comment):
         if not self._client:
             return False
+        self.connect()
         if comment_status := self._client.stories.create_story_for_task(task_gid, {"type": "comment", "text": comment}):
             status = comment_status.get("text", "")
         return len(status) > 0

src/n0s1/controllers/confluence_controller.py

@@ -12,18 +12,18 @@
 class ConfluenceController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
         self._url = None
         self._user = None
         self._password = None
 
-    def set_config(self, config):
+    def set_config(self, config=None):
+        super().set_config(config)
         from atlassian import Confluence
-        SERVER = config.get("server", "")
-        EMAIL = config.get("email", "")
-        TOKEN = config.get("token", "")
-        TIMEOUT = config.get("timeout", -1)
-        VERIFY_SSL = not config.get("insecure", False)
+        SERVER = self._config.get("server", "")
+        EMAIL = self._config.get("email", "")
+        TOKEN = self._config.get("token", "")
+        TIMEOUT = self._config.get("timeout", -1)
+        VERIFY_SSL = not self._config.get("insecure", False)
         self._url = SERVER
         self._user = EMAIL
         self._password = TOKEN
@@ -132,6 +132,7 @@ def get_data(self, include_coments=False, limit=None):
         finished = False
         while not finished:
             try:
+                self.connect()
                 res = self._client.get_all_spaces(start=space_start, limit=limit)
                 spaces = res.get("results", [])
             except Exception as e:
@@ -150,6 +151,7 @@ def get_data(self, include_coments=False, limit=None):
                     pages_finished = False
                     while not pages_finished:
                         try:
+                            self.connect()
                             pages = self._client.get_all_pages_from_space(key, start=pages_start, limit=limit)
                         except ApiPermissionError as e:
                             message = str(e) + f" get_all_pages_from_space({key}, start={pages_start}, limit={limit}). Skipping..."
@@ -169,6 +171,7 @@ def get_data(self, include_coments=False, limit=None):
                             title = p.get("title", "")
                             page_id = p.get("id", "")
                             try:
+                                self.connect()
                                 body = self._client.get_page_by_id(page_id, expand="body.storage")
                             except Exception as e:
                                 message = str(e) + f" get_page_by_id({page_id})"
@@ -184,6 +187,7 @@ def get_data(self, include_coments=False, limit=None):
                                 comments_finished = False
                                 while not comments_finished:
                                     try:
+                                        self.connect()
                                         comments_response = self._client.get_page_comments(page_id, expand="body.storage", start=comments_start, limit=limit)
                                         comments_result = comments_response.get("results", [])
                                     except Exception as e:
@@ -216,6 +220,7 @@ def post_comment(self, issue, comment):
         comment = comment.replace("#", "0")
         comment = html.escape(comment, quote=True)
         status = -1
+        self.connect()
         if comment_status := self._client.add_comment(issue, comment):
             status = comment_status.get("id", "")
         return int(status) > 0

src/n0s1/controllers/hollow_controller.py

@@ -3,18 +3,36 @@
 
 class HollowController:
     def __init__(self):
-        self.client = None
+        self._client = None
+        self._config = None
+        self._check_connection_after = 200
+        self._requests_counter = 0
         self.log_message_callback = None
 
-    def set_config(self, config):
-        return self.is_connected()
+    def set_config(self, config=None):
+        if config:
+            self._config = config
+        return self._config is not None
 
     def set_log_message_callback(self, log_message_callback):
         self.log_message_callback = log_message_callback
 
     def get_name(self):
         return "Hollow"
 
+    def connect(self):
+        self._requests_counter += 1
+        if self._requests_counter > self._check_connection_after:
+            self._requests_counter = 0
+            if not self.is_connected():
+                # Force new connection
+                self._client = None
+
+        if not self._client:
+            if self.set_config():
+                return self.is_connected()
+        return True
+
     def is_connected(self):
         return False
 

src/n0s1/controllers/jira_controller.py

@@ -10,15 +10,15 @@
 class JiraController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
 
-    def set_config(self, config):
+    def set_config(self, config=None):
+        super().set_config(config)
         from jira import JIRA
-        SERVER = config.get("server", "")
-        EMAIL = config.get("email", "")
-        TOKEN = config.get("token", "")
-        TIMEOUT = config.get("timeout", -1)
-        VERIFY_SSL = not config.get("insecure", False)
+        SERVER = self._config.get("server", "")
+        EMAIL = self._config.get("email", "")
+        TOKEN = self._config.get("token", "")
+        TIMEOUT = self._config.get("timeout", -1)
+        VERIFY_SSL = not self._config.get("insecure", False)
         options = {"verify": VERIFY_SSL}
         if EMAIL and len(EMAIL) > 0:
             if TIMEOUT and TIMEOUT > 0:
@@ -72,6 +72,7 @@ def get_data(self, include_coments=False, limit=None):
         if not limit or limit < 0:
             limit = 50
         try:
+            self.connect()
             projects = self._client.projects()
         except Exception as e:
             message = str(e) + f" client.projects()"
@@ -85,6 +86,7 @@ def get_data(self, include_coments=False, limit=None):
             issue_start = start
             while not issues_finished:
                 try:
+                    self.connect()
                     issues = self._client.search_issues(ql, startAt=issue_start, maxResults=limit)
                 except JIRAError as e:
                     self.log_message(f"Error while searching issues on Jira project: [{key}]. Skipping...", logging.WARNING)
@@ -106,6 +108,7 @@ def get_data(self, include_coments=False, limit=None):
                     comments = []
                     if include_coments:
                         try:
+                            self.connect()
                             issue_comments = self._client.comments(issue.id)
                             comments.extend(c.body for c in issue_comments)
                         except Exception as e:
@@ -121,6 +124,7 @@ def post_comment(self, issue, comment):
         if not self._client:
             return False
         comment = comment.replace("#", "0")
+        self.connect()
         comment_status = self._client.add_comment(issue, body=comment)
         status = comment_status.id
         return bool(status and len(status) > 0 and int(status) > 0)

src/n0s1/controllers/linear_controller.py

@@ -14,10 +14,10 @@
 class LinearController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
 
-    def set_config(self, config):
-        TOKEN = config.get("token", "")
+    def set_config(self, config=None):
+        super().set_config(config)
+        TOKEN = self._config.get("token", "")
         headers = {
             "Content-Type": "application/json",
             "Authorization": TOKEN,
@@ -55,7 +55,9 @@ def is_connected(self):
     def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return {}
+        self.connect()
         for linear_data in self._client.get_issues_and_comments(20):
+            self.connect()
             for edge in linear_data.get("data", {}).get("issues", {}).get("edges", []):
                 item = edge.get("node", {})
                 url = item.get("url", "")

src/n0s1/controllers/slack_controller.py

@@ -12,12 +12,12 @@
 class SlackController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
 
-    def set_config(self, config):
+    def set_config(self, config=None):
+        super().set_config(config)
         from slack_sdk import WebClient
         from slack_sdk.errors import SlackApiError
-        TOKEN = config.get("token", "")
+        TOKEN = self._config.get("token", "")
         self._client = WebClient(token=TOKEN)
         return self.is_connected()
 
@@ -75,6 +75,7 @@ def post_comment(self, issue, comment):
         try:
             channel_id, thread_ts = self.extract_channel_id_and_ts(issue)
             if comment and len(comment) > 0 and len(channel_id) > 0 and len(thread_ts) > 0:
+                self.connect()
                 response = self._client.chat_postMessage(
                     channel=channel_id,
                     text=comment,
@@ -121,6 +122,7 @@ def search_with_rate_limit(self, query, sort, cursor):
         from slack_sdk.errors import SlackApiError
         response = None
         try:
+            self.connect()
             response = self._client.search_messages(query=query, sort=sort, cursor=cursor)
         except SlackApiError as ex:
             message = str(ex) + f" client.search_messages()"

src/n0s1/controllers/wrike_controller.py

@@ -8,13 +8,14 @@
 except Exception:
     import n0s1.controllers.hollow_controller as hollow_controller
 
+
 class WrikeController(hollow_controller.HollowController):
     def __init__(self):
         super().__init__()
-        self._client = None
 
-    def set_config(self, config):
-        TOKEN = config.get("token", "")
+    def set_config(self, config=None):
+        super().set_config(config)
+        TOKEN = self._config.get("token", "")
         base_url = "https://www.wrike.com/api/v4"
         self._client = Wrike(base_url, TOKEN)
         return self.is_connected()
@@ -61,6 +62,7 @@ def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return {}
 
+        self.connect()
         t = Tasks(self._client, parameters={"fields": ["description"]})
         response = t.query__tasks()
         tasks = {}
@@ -78,6 +80,7 @@ def get_data(self, include_coments=False, limit=None):
             comments = []
             if task_id := t.get("id", None):
                 if include_coments:
+                    self.connect()
                     comments_obj = Comments(self._client, [task_id])
                     response = comments_obj.query__tasks_taskId_comments()
                     json_data = {}
@@ -98,6 +101,7 @@ def post_comment(self, task_id, comment):
             return False
         comment = comment.replace("<REDACTED>", "**********")
         comment = comment.replace("\n", "<br>")
+        self.connect()
         comments_obj = Comments(self._client, [task_id], parameters={"text": comment, "plainText": False})
         if comments_obj:
             response = comments_obj.create__tasks_taskId_comments()