From 5ab90c3e5fb1764d5ca6062c02f9578a7d5a9c7b Mon Sep 17 00:00:00 2001
From: 0x41424142 <89279733+0x41424142@users.noreply.github.com>
Date: Tue, 23 Jul 2024 22:15:30 -0400
Subject: [PATCH] Added chunk_count to hld, fixed bs4 warning suppression

---
 qualyspy/vmdr/data_classes/detection.py   |  2 +-
 qualyspy/vmdr/data_classes/kb_entry.py    |  2 +-
 qualyspy/vmdr/get_host_list_detections.py | 44 +++++++++++++++++------
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/qualyspy/vmdr/data_classes/detection.py b/qualyspy/vmdr/data_classes/detection.py
index c26ef7f..3c3a0cd 100644
--- a/qualyspy/vmdr/data_classes/detection.py
+++ b/qualyspy/vmdr/data_classes/detection.py
@@ -143,7 +143,7 @@ def __post_init__(self):
 
         # clean up fields that have html tags
         with catch_warnings():
-            simplefilter("ignore") # ignore the warning about the html.parser
+            simplefilter("ignore")  # ignore the warning about the html.parser
             for field in HTML_FIELDS:
                 setattr(
                     self,
diff --git a/qualyspy/vmdr/data_classes/kb_entry.py b/qualyspy/vmdr/data_classes/kb_entry.py
index 81fe049..69ad5ab 100644
--- a/qualyspy/vmdr/data_classes/kb_entry.py
+++ b/qualyspy/vmdr/data_classes/kb_entry.py
@@ -199,7 +199,7 @@ def __post_init__(self):
                 setattr(self, field, bool(getattr(self, field)))
 
         with catch_warnings():
-            simplefilter("ignore") # ignore the warning about the html.parser
+            simplefilter("ignore")  # ignore the warning about the html.parser
             for field in HTML_FIELDS:
                 if getattr(self, field) is not None:
                     setattr(
diff --git a/qualyspy/vmdr/get_host_list_detections.py b/qualyspy/vmdr/get_host_list_detections.py
index 0ee3a8c..c7d7f90 100644
--- a/qualyspy/vmdr/get_host_list_detections.py
+++ b/qualyspy/vmdr/get_host_list_detections.py
@@ -22,6 +22,7 @@
 
 LOCK = Lock()
 
+
 def normalize_id_list(id_list):
     """
     normalize_id_list - formats the kwarg ids, if it is passed.
@@ -286,9 +287,11 @@ def create_id_queue(
     for i in range(0, len(id_list), chunk_size):
         id_queue.put(id_list[i : i + chunk_size])
 
-    singular_chunk = True if id_queue.qsize() == 1 else False 
+    singular_chunk = True if id_queue.qsize() == 1 else False
 
-    print(f"Queue created with {id_queue.qsize()} {'chunks' if not singular_chunk else 'chunk'} of ~{chunk_size} IDs{' each.' if not singular_chunk else '.'}")
+    print(
+        f"Queue created with {id_queue.qsize()} {'chunks' if not singular_chunk else 'chunk'} of ~{chunk_size} IDs{' each.' if not singular_chunk else '.'}"
+    )
 
     return id_queue
 
@@ -317,10 +320,19 @@ def get_hld(
         BaseList: A list of VMDRHost objects, with their DETECTIONS attribute populated.
     """
 
-    #Ensure that threads, chunk_size, and page_count (if not 'all') are all integers above 0
-    if any([threads < 1, chunk_size < 1, (page_count != "all" and page_count < 1), (chunk_count != "all" and chunk_count < 1)]):
-        raise ValueError("threads, chunk_size, page_count (if not 'all') and chunk_count (if not 'all') must all be integers above 0.")
-    
+    # Ensure that threads, chunk_size, and page_count (if not 'all') are all integers above 0
+    if any(
+        [
+            threads < 1,
+            chunk_size < 1,
+            (page_count != "all" and page_count < 1),
+            (chunk_count != "all" and chunk_count < 1),
+        ]
+    ):
+        raise ValueError(
+            "threads, chunk_size, page_count (if not 'all') and chunk_count (if not 'all') must all be integers above 0."
+        )
+
     # Make sure the user hasn't set threads to more than the cpu count
     if threads > cpu_count():
         print(
@@ -335,8 +347,12 @@ def get_hld(
         )
         threads = rl["X-Concurrency-Limit-Limit"]
 
-    print(f"Pulling/creating queue for full ID list...") if not kwargs.get("ids") else print(
-        f"Pulling/creating queue for user-specified IDs: {kwargs.get('ids')}..."
+    (
+        print(f"Pulling/creating queue for full ID list...")
+        if not kwargs.get("ids")
+        else print(
+            f"Pulling/creating queue for user-specified IDs: {kwargs.get('ids')}..."
+        )
     )
 
     id_queue = create_id_queue(auth, chunk_size=chunk_size, ids=kwargs.get("ids"))
@@ -384,7 +400,9 @@ def threaded_hld_worker(
         pages_pulled = 0
         chunks_pulled = 0
         try:
-            ids = id_queue.get_nowait() #nowait allows us to check if the queue is empty without blocking
+            ids = (
+                id_queue.get_nowait()
+            )  # nowait allows us to check if the queue is empty without blocking
         except Empty:
             with LOCK:
                 print(f"{current_thread().name} - Queue is empty. Terminating thread.")
@@ -409,9 +427,13 @@ def threaded_hld_worker(
             break
         if pages_pulled == page_count:
             with LOCK:
-                print(f"{current_thread().name} - Thread has pulled all pages. Terminating thread.")
+                print(
+                    f"{current_thread().name} - Thread has pulled all pages. Terminating thread."
+                )
             break
         if chunks_pulled == chunk_count:
             with LOCK:
-                print(f"{current_thread().name} - Thread has pulled all chunks. Terminating thread.")
+                print(
+                    f"{current_thread().name} - Thread has pulled all chunks. Terminating thread."
+                )
             break