Merge pull request #16 from spark1security/bp-timeout

Added timeout and limit to Jira and Confluence

src/n0s1/__init__.py

@@ -1 +1 @@
-__version__ = "1.0.17"
+__version__ = "1.0.18"

src/n0s1/config/config.yaml

@@ -7,6 +7,8 @@ general_params:
   skip_comment: false
   show_matched_secret_on_logs: false
   report_format: "n0s1"
+  timeout: null
+  limit: null
 
 comment_params:
   bot_name: "n0s1 bot"

src/n0s1/controllers/asana_controller.py

@@ -1,8 +1,15 @@
 import logging
 
 
-class AsanaControler():
+try:
+    from . import hollow_controller as hollow_controller
+except Exception:
+    import n0s1.controllers.hollow_controller as hollow_controller
+
+
+class AsanaControler(hollow_controller.HollowController):
     def __init__(self):
+        super().__init__()
         self._client = None
 
     def set_config(self, config):
@@ -17,9 +24,9 @@ def get_name(self):
     def is_connected(self):
         if self._client:
             if user := self._client.users.get_user("me"):
-                logging.info(f"Logged to {self.get_name()} as {user}")
+                self.log_message(f"Logged to {self.get_name()} as {user}")
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
                 return False
 
             if spaces := self._client.workspaces.get_workspaces():
@@ -37,14 +44,14 @@ def is_connected(self):
                     if project_found:
                         return True
                     else:
-                        logging.error(f"Unable to list {self.get_name()} projects. Check your permissions.")
+                        self.log_message(f"Unable to list {self.get_name()} projects. Check your permissions.", logging.ERROR)
                 else:
-                    logging.error(f"Unable to list {self.get_name()} workspaces. Check your permissions.")
+                    self.log_message(f"Unable to list {self.get_name()} workspaces. Check your permissions.", logging.ERROR)
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
         return False
 
-    def get_data(self, include_coments=False):
+    def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return None, None, None, None, None
 

src/n0s1/controllers/confluence_controller.py

@@ -1,9 +1,17 @@
 import html
 import logging
+import time
 
 
-class ConfluenceControler():
+try:
+    from . import hollow_controller as hollow_controller
+except Exception:
+    import n0s1.controllers.hollow_controller as hollow_controller
+
+
+class ConfluenceControler(hollow_controller.HollowController):
     def __init__(self):
+        super().__init__()
         self._client = None
         self._url = None
         self._user = None
@@ -14,13 +22,20 @@ def set_config(self, config):
         SERVER = config.get("server", "")
         EMAIL = config.get("email", "")
         TOKEN = config.get("token", "")
+        TIMEOUT = config.get("timeout", -1)
         self._url = SERVER
         self._user = EMAIL
         self._password = TOKEN
         if EMAIL and len(EMAIL) > 0:
-            self._client = Confluence(url=SERVER, username=EMAIL, password=TOKEN)
+            if TIMEOUT and TIMEOUT > 0:
+                self._client = Confluence(url=SERVER, username=EMAIL, password=TOKEN, timeout=TIMEOUT)
+            else:
+                self._client = Confluence(url=SERVER, username=EMAIL, password=TOKEN)
         else:
-            self._client = Confluence(url=SERVER, token=TOKEN)
+            if TIMEOUT and TIMEOUT > 0:
+                self._client = Confluence(url=SERVER, token=TOKEN, timeout=TIMEOUT)
+            else:
+                self._client = Confluence(url=SERVER, token=TOKEN)
         return self.is_connected()
 
     def get_name(self):
@@ -53,7 +68,7 @@ def _get_request(self, url):
                     headers=headers
                 )
         except Exception as e:
-            logging.info(e)
+            self.log_message(str(e))
         return response
 
     def get_current_user(self):
@@ -78,9 +93,9 @@ def get_current_user(self):
     def is_connected(self):
         if self._client:
             if user := self.get_current_user():
-                logging.info(f"Logged to {self.get_name()} as {user}")
+                self.log_message(f"Logged to {self.get_name()} as {user}")
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
                 return False
 
             if spaces := self._client.get_all_spaces():
@@ -98,56 +113,79 @@ def is_connected(self):
                     if page_found:
                         return True
                     else:
-                        logging.error(f"Unable to list {self.get_name()} pages. Check your permissions.")
+                        self.log_message(f"Unable to list {self.get_name()} pages. Check your permissions.", logging.ERROR)
                 else:
-                    logging.error(f"Unable to list {self.get_name()} spaces. Check your permissions.")
+                    self.log_message(f"Unable to list {self.get_name()} spaces. Check your permissions.", logging.ERROR)
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
         return False
 
-    def get_data(self, include_coments=False):
+    def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return None, None, None, None, None
 
-        start = 0
-        limit = 500
+        space_start = 0
+        if not limit or limit < 0:
+            limit = 50
         finished = False
         while not finished:
-            res = self._client.get_all_spaces(start=start, limit=limit)
-            start = limit
-            limit += start
-            spaces = res.get("results", [])
+            try:
+                res = self._client.get_all_spaces(start=space_start, limit=limit)
+                spaces = res.get("results", [])
+            except Exception as e:
+                message = str(e) + f" get_all_spaces(start={space_start}, limit={limit})"
+                self.log_message(message, logging.WARNING)
+                spaces = [{}]
+                time.sleep(1)
+                continue
+            space_start += limit
 
             for s in spaces:
                 key = s.get("key", "")
-                logging.info(f"Scanning Confluence space: [{key}]...")
+                self.log_message(f"Scanning Confluence space: [{key}]...")
                 if len(key) > 0:
                     pages_start = 0
-                    pages_limit = 50
                     pages_finished = False
                     while not pages_finished:
-                        pages = self._client.get_all_pages_from_space(key, start=pages_start, limit=pages_limit)
-                        pages_start = pages_limit
-                        pages_limit += pages_start
+                        try:
+                            pages = self._client.get_all_pages_from_space(key, start=pages_start, limit=limit)
+                        except Exception as e:
+                            message = str(e) + f" get_all_pages_from_space({key}, start={pages_start}, limit={limit})"
+                            self.log_message(message, logging.WARNING)
+                            pages = [{}]
+                            time.sleep(1)
+                            continue
+                        pages_start += limit
 
                         for p in pages:
                             comments = []
                             title = p.get("title", "")
                             page_id = p.get("id", "")
-                            body = self._client.get_page_by_id(page_id, expand="body.storage")
+                            try:
+                                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})"
+                                self.log_message(message, logging.WARNING)
+                                body = {}
+                                time.sleep(1)
+                                continue
+
                             description = body.get("body", {}).get("storage", {}).get("value", "")
-                            url = body.get("_links", {}).get("base") + p.get("_links", {}).get("webui", "")
-                            if include_coments:
+                            url = body.get("_links", {}).get("base", "") + p.get("_links", {}).get("webui", "")
+                            if len(page_id) > 0 and include_coments:
                                 comments_start = 0
-                                comments_limit = 25
                                 comments_finished = False
                                 while not comments_finished:
-                                    comments_response = self._client.get_page_comments(page_id, expand="body.storage",
-                                                                                       start=comments_start,
-                                                                                       limit=comments_limit)
-                                    comments_start = comments_limit
-                                    comments_limit += comments_start
-                                    comments_result = comments_response.get("results", [])
+                                    try:
+                                        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:
+                                        message = str(e) + f" get_page_comments({page_id}, expand=\"body.storage\", start={comments_start}, limit={limit})"
+                                        self.log_message(message, logging.WARNING)
+                                        comments_result = [{}]
+                                        time.sleep(1)
+                                        continue
+                                    comments_start += limit
 
                                     for c in comments_result:
                                         comment = c.get("body", {}).get("storage", {}).get("value", "")

src/n0s1/controllers/hollow_controller.py

@@ -0,0 +1,31 @@
+import logging
+
+
+class HollowController:
+    def __init__(self):
+        self.client = None
+        self.log_message_callback = None
+
+    def set_config(self, config):
+        return self.is_connected()
+
+    def set_log_message_callback(self, log_message_callback):
+        self.log_message_callback = log_message_callback
+
+    def get_name(self):
+        return "Hollow"
+
+    def is_connected(self):
+        return False
+
+    def get_data(self, include_coments=False, limit=None):
+        return None, None, None, None, None
+
+    def post_comment(self, issue, comment):
+        return self.is_connected()
+
+    def log_message(self, message, level=logging.INFO):
+        if self.log_message_callback:
+            self.log_message_callback(message, level)
+        else:
+            print(message)

src/n0s1/controllers/jira_controller.py

@@ -1,19 +1,33 @@
 import logging
+import time
 
+try:
+    from . import hollow_controller as hollow_controller
+except Exception:
+    import n0s1.controllers.hollow_controller as hollow_controller
 
-class JiraControler():
+
+class JiraControler(hollow_controller.HollowController):
     def __init__(self):
+        super().__init__()
         self._client = None
 
     def set_config(self, config):
         from jira import JIRA
         SERVER = config.get("server", "")
         EMAIL = config.get("email", "")
         TOKEN = config.get("token", "")
+        TIMEOUT = config.get("timeout", -1)
         if EMAIL and len(EMAIL) > 0:
-            self._client = JIRA(SERVER, basic_auth=(EMAIL, TOKEN))
+            if TIMEOUT and TIMEOUT > 0:
+                self._client = JIRA(SERVER, basic_auth=(EMAIL, TOKEN), timeout=TIMEOUT)
+            else:
+                self._client = JIRA(SERVER, basic_auth=(EMAIL, TOKEN))
         else:
-            self._client = JIRA(SERVER, token_auth=TOKEN)
+            if TIMEOUT and TIMEOUT > 0:
+                self._client = JIRA(SERVER, token_auth=TOKEN, timeout=TIMEOUT)
+            else:
+                self._client = JIRA(SERVER, token_auth=TOKEN)
         return self.is_connected()
 
     def get_name(self):
@@ -22,9 +36,9 @@ def get_name(self):
     def is_connected(self):
         if self._client:
             if user := self._client.myself():
-                logging.info(f"Logged to {self.get_name()} as {user}")
+                self.log_message(f"Logged to {self.get_name()} as {user}")
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
                 return False
 
             if projects := self._client.projects():
@@ -41,28 +55,58 @@ def is_connected(self):
                     if issue_found:
                         return True
                     else:
-                        logging.error(f"Unable to list {self.get_name()} issues. Check your permissions.")
+                        self.log_message(f"Unable to list {self.get_name()} issues. Check your permissions.", logging.ERROR)
                 else:
-                    logging.error(f"Unable to list {self.get_name()} projects. Check your permissions.")
+                    self.log_message(f"Unable to list {self.get_name()} projects. Check your permissions.", logging.ERROR)
             else:
-                logging.error(f"Unable to connect to {self.get_name()} instance. Check your credentials.")
+                self.log_message(f"Unable to connect to {self.get_name()} instance. Check your credentials.", logging.ERROR)
         return False
 
-    def get_data(self, include_coments=False):
+    def get_data(self, include_coments=False, limit=None):
         if not self._client:
             return None, None, None, None, None
-        for key in self._client.projects():
+        start = 0
+        if not limit or limit < 0:
+            limit = 50
+        try:
+            projects = self._client.projects()
+        except Exception as e:
+            message = str(e) + f" client.projects()"
+            self.log_message(message, logging.WARNING)
+            projects = []
+
+        for key in projects:
             ql = f"project = '{key}'"
-            logging.info(f"Scanning Jira project: [{key}]...")
-            for issue in self._client.search_issues(ql):
-                url = issue.self.split('/rest/api')[0] + "/browse/" + issue.key;
-                title = issue.fields.summary
-                description = issue.fields.description
-                comments = []
-                if include_coments:
-                    issue_comments = self._client.comments(issue.id)
-                    comments.extend(c.body for c in issue_comments)
-                yield title, description, comments, url, issue.key
+            self.log_message(f"Scanning Jira project: [{key}]...")
+            issues_finished = False
+            issue_start = start
+            while not issues_finished:
+                try:
+                    issues = self._client.search_issues(ql, startAt=issue_start, maxResults=limit)
+                except Exception as e:
+                    message = str(e) + f" client.search_issues({ql}, startAt={issue_start}, maxResults={limit})"
+                    self.log_message(message, logging.WARNING)
+                    issues = [{}]
+                    time.sleep(1)
+                    continue
+                issue_start += limit
+                issues_finished = len(issues) <= 0
+                for issue in issues:
+                    url = issue.self.split('/rest/api')[0] + "/browse/" + issue.key;
+                    title = issue.fields.summary
+                    description = issue.fields.description
+                    comments = []
+                    if include_coments:
+                        try:
+                            issue_comments = self._client.comments(issue.id)
+                            comments.extend(c.body for c in issue_comments)
+                        except Exception as e:
+                            message = str(e) + f" client.comments({issue.id})"
+                            self.log_message(message, logging.WARNING)
+                            comments = []
+                            time.sleep(1)
+
+                    yield title, description, comments, url, issue.key
 
     def post_comment(self, issue, comment):
         if not self._client: