Fix Automaton.graph() with indirection

scapy/automaton.py

@@ -783,12 +783,19 @@ def build_graph(self):
                         [("red", k, v) for k, v in self.recv_conditions.items()] +  # noqa: E501
                         [("orange", k, v) for k, v in self.ioevents.items()]):
             for f in v:
-                for n in f.__code__.co_names + f.__code__.co_consts:
+                names = list(f.__code__.co_names + f.__code__.co_consts)
+                while names:
+                    n = names.pop()
                     if n in self.states:
                         line = f.atmt_condname
                         for x in self.actions[f.atmt_condname]:
                             line += "\\l>[%s]" % x.__name__
                         s += '\t"%s" -> "%s" [label="%s", color=%s];\n' % (k, n, line, c)  # noqa: E501
+                    elif n in self.__dict__:
+                        # test for function
+                        if callable(self.__dict__[n]):
+                            names.extend(self.__dict__[n].__code__.co_names)
+                            names.extend(self.__dict__[n].__code__.co_consts)
         for k, timers in self.timeout.items():
             for timer in timers:
                 for n in (timer._func.__code__.co_names +

test/scapy/automaton.uts

@@ -446,6 +446,27 @@ graph = HelloWorld.build_graph()
 assert graph.startswith("digraph")
 assert '"BEGIN" -> "END"' in graph
 
+= Automaton graph - with indirection
+~ automaton
+
+class HelloWorld(Automaton):
+    @ATMT.state(initial=1)
+    def BEGIN(self):
+        self.count1 = 0
+        self.count2 = 0
+    @ATMT.condition(BEGIN)
+    def cnd_1(self):
+        self.cnd_generic()
+    def cnd_generic(self):
+        raise END
+    @ATMT.state(final=1)
+    def END(self):
+        pass
+
+graph = HelloWorld.build_graph()
+assert graph.startswith("digraph")
+assert '"BEGIN" -> "END"' in graph
+
 = TCP_client automaton
 ~ automaton netaccess needs_root
 * This test retries on failure because it may fail quite easily