From a6bbf949aa82ad9f7eaec0b5deeb012870e770da Mon Sep 17 00:00:00 2001
From: gpotter2 <10530980+gpotter2@users.noreply.github.com>
Date: Wed, 13 Sep 2023 12:29:01 +0000
Subject: [PATCH] Fix Automaton.graph() with indirection

---
 scapy/automaton.py       |  9 ++++++++-
 test/scapy/automaton.uts | 21 +++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/scapy/automaton.py b/scapy/automaton.py
index a12fb6953c5..5f7b82c23eb 100644
--- a/scapy/automaton.py
+++ b/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 +
diff --git a/test/scapy/automaton.uts b/test/scapy/automaton.uts
index bc103f3126d..ab107b8bc74 100644
--- a/test/scapy/automaton.uts
+++ b/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