Added Local List Check

TepidJesus · May 31, 2023 · f7d9095 · f7d9095
1 parent a8e9968
commit f7d9095
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/job_handler.py b/job_handler.py
@@ -36,6 +36,7 @@ class JobHandler:
     def __init__(self, aws_controller, mode, debug=False):
         self.aws_controller = aws_controller
         self.debug = debug
+        self.stored_wordlists = []
         if mode == "client":
             self.outbound_queue = self.aws_controller.create_queue('deliveryQueue')
             self.control_queue = self.aws_controller.create_queue('controlQueue')
@@ -67,15 +68,17 @@ def send_job(self, job):
                     print(f"[DEBUG] Creating new bucket for job {job.job_id}.")
                 self.wordlist_bucket_name = self.aws_controller.create_bucket("wordlist-bucket")
             file_name = self.get_file_name(job.required_info)
-            response = self.aws_controller.upload_file(job.required_info, self.wordlist_bucket_name, file_name) #TODO: Prevent duplicate uploads
-            if response == False:
-                if self.debug:
-                    print(f"[DEBUG] Failed to Upload Wordlist for job #{job.job_id}. Continuing...")
-                else:
-                    print(f"Error: Failed to create bucket for job {job.job_id}. Continuing...")
-                    print("Please check your AWS S3 Permissions and try again.")
-                job.job_status = STATUS.FAILED
-                return
+            if file_name not in self.stored_wordlists:
+                if not self.aws_controller.upload_file(job.required_info, self.wordlist_bucket_name, file_name):
+                    if self.debug:
+                        print(f"[DEBUG] Failed to Upload Wordlist for job #{job.job_id}. Continuing...")
+                    else:
+                        print(f"Error: Failed to create bucket for job {job.job_id}. Continuing...")
+                        print("Please check your AWS S3 Permissions and try again.")
+                    job.job_status = STATUS.FAILED
+                    return
+            elif self.debug:
+                print(f"[DEBUG] Wordlist {file_name} already Uploaded. Continuing...")
 
             job.required_info = (file_name, self.wordlist_bucket_name)