Consider nxdomain_ancestor and follow_mx

1 files changed, 10 insertions, 0 deletions

diff --git a/dnsviz/commands/probe.py b/dnsviz/commands/probe.py
index 3e1c3b3..a10577e 100644
--- a/dnsviz/commands/probe.py
+++ b/dnsviz/commands/probe.py
@@ -195,6 +195,8 @@ class MultiProcessAnalystMixin(object):
 
         if name_obj.parent is not None:
             self.refresh_dependency_references(name_obj.parent, trace+[name_obj.name])
+        if name_obj.nxdomain_ancestor is not None:
+            self.refresh_dependency_references(name_obj.nxdomain_ancestor, trace+[name_obj.name])
         if name_obj.dlv_parent is not None:
             self.refresh_dependency_references(name_obj.dlv_parent, trace+[name_obj.name])
 
@@ -222,6 +224,14 @@ class MultiProcessAnalystMixin(object):
                     except KeyError:
                         time.sleep(1)
                 self.refresh_dependency_references(name_obj.ns_dependencies[ns], trace+[name_obj.name])
+        if self.follow_mx:
+            for target in name_obj.mx_targets:
+                while name_obj.mx_targets[target] is None:
+                    try:
+                        name_obj.mx_targets[target] = self.analysis_cache[target]
+                    except KeyError:
+                        time.sleep(1)
+                self.refresh_dependency_references(name_obj.mx_targets[target], trace+[name_obj.name])
 
     def analyze(self):
         name_obj = super(MultiProcessAnalystMixin, self).analyze()