yeahhh

artemis/modules/domain_scanner.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 from karton.core import Task
 from artemis.resolvers import lookup, ResolutionException
-
 from artemis.binds import TaskStatus, TaskType
 from artemis.module_base import ArtemisBase
 
@@ -13,36 +12,9 @@ class DomainScanner(ArtemisBase):
     identity = "domain_scanner"
     filters = [
         {"type": TaskType.DOMAIN.value},
+        {"type": TaskType.NEW.value},
     ]
 
-    def run(self, current_task: Task) -> None:
-        domain = current_task.get_payload(TaskType.DOMAIN)
-        domain_exists = self.check_domain_exists(domain)
-
-        if domain_exists:
-            self.db.save_task_result(
-                task=current_task,
-                status=TaskStatus.OK,
-                data={
-                    "domain": domain,
-                    "exists": True
-                }
-            )
-            # Create a new task for the next module
-            new_task = Task(
-                {"type": TaskType.DOMAIN},
-                payload={TaskType.DOMAIN: domain},
-                payload_persistent={"domain_exists": True}
-            )
-            self.add_task(current_task, new_task)
-        else:
-            self.db.save_task_result(
-                task=current_task,
-                status=TaskStatus.ERROR,
-                status_reason=f"Domain does not exist: {domain}",
-                data={"domain": domain, "exists": False}
-            )
-
     def check_domain_exists(self, domain: str) -> bool:
         try:
             # Check for NS records
@@ -56,5 +28,43 @@ def check_domain_exists(self, domain: str) -> bool:
         except ResolutionException:
             return False
 
+    def run(self, current_task: Task) -> None:
+        domains = current_task.get_payload(TaskType.DOMAIN)
+        if isinstance(domains, str):
+            domains = [domains]
+
+        existing_domains = []
+        non_existing_domains = []
+
+        for domain in domains:
+            if self.check_domain_exists(domain):
+                existing_domains.append(domain)
+            else:
+                non_existing_domains.append(domain)
+
+        self.db.save_task_result(
+            task=current_task,
+            status=TaskStatus.OK,
+            data={
+                "existing_domains": existing_domains,
+                "non_existing_domains": non_existing_domains
+            }
+        )
+
+        # Create a new task for existing domains
+        if existing_domains:
+            self.add_task(Task(
+                headers={
+                    "type": TaskType.DOMAIN.value,
+                },
+                payload={
+                    TaskType.DOMAIN.value: existing_domains
+                }
+            ))
+
+        # Log or handle non-existing domains
+        if non_existing_domains:
+            self.log.info(f"The following domains do not exist and will not be scanned: {non_existing_domains}")
+
 if __name__ == "__main__":
     DomainScanner().loop()