Repository: secdev/scapy

Issue:
[Enhancement] Inflexible Automaton timeout configuration
### Brief description

The current implementation of `Automaton_metaclass` creates a `timeout` list in `__new__`, which means that timeout values are set during code sourcing and can only be changed by modifying `self.timeout` manually. Meanwhile, many protocols (e.g. ERPS) require timeout configuration. This could be supported by allowing functions as `ATMT.timeout` arguments and expanding them during `__init__`. 
I'll try to implement it if this feature is approved.

### Scapy version

6d7184e8bec5102dfa66bcc10432a30a7e0dcf3a

### Python version

3.10

### Operating system

All

### Additional environment information

_No response_

### How to reproduce

-

### Actual result

_No response_

### Expected result

_No response_

### Related resources

_No response_


Discussion:
Hi !

Cant you just dynamically create the python class? For instance

```
def create(timeout):
    class A(Automaton):
        [...]
        @ATMT.timeout(timeout)
        def tmout(self):
            pass
```
Maybe a timeout class would be better? Then lazy timeout sort could be done, allowing runtime reconfiguration.
This does sound a bit like bloat. Feel free to try to prototype something.. It's hard to see an actual useful case

Resolving patch:
diff --git a/scapy/automaton.py b/scapy/automaton.py
--- a/scapy/automaton.py
+++ b/scapy/automaton.py
@@ -258,6 +258,124 @@ def __repr__(self):
                                          if not k.startswith("_"))
 
 
+class Timer():
+    def __init__(self, time, prio=0, autoreload=False):
+        # type: (Union[int, float], int, bool) -> None
+        self._timeout = float(time)  # type: float
+        self._time = 0  # type: float
+        self._just_expired = True
+        self._expired = True
+        self._prio = prio
+        self._func = _StateWrapper()
+        self._autoreload = autoreload
+
+    def get(self):
+        # type: () -> float
+        return self._timeout
+
+    def set(self, val):
+        # type: (float) -> None
+        self._timeout = val
+
+    def _reset(self):
+        # type: () -> None
+        self._time = self._timeout
+        self._expired = False
+        self._just_expired = False
+
+    def _reset_just_expired(self):
+        # type: () -> None
+        self._just_expired = False
+
+    def _running(self):
+        # type: () -> bool
+        return self._time > 0
+
+    def _remaining(self):
+        # type: () -> float
+        return max(self._time, 0)
+
+    def _decrement(self, time):
+        # type: (float) -> None
+        self._time -= time
+        if self._time <= 0:
+            if not self._expired:
+                self._just_expired = True
+                if self._autoreload:
+                    # take overshoot into account
+                    self._time = self._timeout + self._time
+                else:
+                    self._expired = True
+                    self._time = 0
+
+    def __lt__(self, obj):
+        # type: (Timer) -> bool
+        return ((self._time < obj._time) if self._time != obj._time
+                else (self._prio < obj._prio))
+
+    def __gt__(self, obj):
+        # type: (Timer) -> bool
+        return ((self._time > obj._time) if self._time != obj._time
+                else (self._prio > obj._prio))
+
+    def __eq__(self, obj):
+        # type: (Any) -> bool
+        if not isinstance(obj, Timer):
+            raise NotImplementedError()
+        return (self._time == obj._time) and (self._prio == obj._prio)
+
+    def __repr__(self):
+        # type: () -> str
+        return "<Timer %f(%f)>" % (self._time, self._timeout)
+
+
+class _TimerList():
+    def __init__(self):
+        # type: () -> None
+        self.timers = []  # type: list[Timer]
+
+    def add_timer(self, timer):
+        # type: (Timer) -> None
+        self.timers.append(timer)
+
+    def reset(self):
+        # type: () -> None
+        for t in self.timers:
+            t._reset()
+
+    def decrement(self, time):
+        # type: (float) -> None
+        for t in self.timers:
+            t._decrement(time)
+
+    def expired(self):
+        # type: () -> list[Timer]
+        lst = [t for t in self.timers if t._just_expired]
+        lst.sort(key=lambda x: x._prio, reverse=True)
+        for t in lst:
+            t._reset_just_expired()
+        return lst
+
+    def until_next(self):
+        # type: () -> float
+        try:
+            return min([t._remaining() for t in self.timers if t._running()])
+        except ValueError:
+            return 0
+
+    def count(self):
+        # type: () -> int
+        return len(self.timers)
+
+    def __iter__(self):
+        # type: () -> Iterator[Timer]
+        return self.timers.__iter__()
+
+    def __repr__(self):
+        # type: () -> str
+        return self.timers.__repr__()
+
+
 class _instance_state:
     def __init__(self, instance):
         # type: (Any) -> None
@@ -307,7 +425,7 @@ class _StateWrapper:
     atmt_as_supersocket = None  # type: Optional[str]
     atmt_condname = None        # type: str
     atmt_ioname = None          # type: str
-    atmt_timeout = None         # type: int
+    atmt_timeout = None         # type: Timer
     atmt_cond = None            # type: Dict[str, int]
     __code__ = None             # type: types.CodeType
     __call__ = None             # type: Callable[..., ATMT.NewStateRequested]
@@ -438,12 +556,26 @@ def deco(f, state=state):
 
     @staticmethod
     def timeout(state, timeout):
-        # type: (_StateWrapper, int) -> Callable[[_StateWrapper, _StateWrapper, int], _StateWrapper]  # noqa: E501
-        def deco(f, state=state, timeout=timeout):
-            # type: (_StateWrapper, _StateWrapper, int) -> _StateWrapper
+        # type: (_StateWrapper, Union[int, float]) -> Callable[[_StateWrapper, _StateWrapper, Timer], _StateWrapper]  # noqa: E501
+        def deco(f, state=state, timeout=Timer(timeout)):
+            # type: (_StateWrapper, _StateWrapper, Timer) -> _StateWrapper
+            f.atmt_type = ATMT.TIMEOUT
+            f.atmt_state = state.atmt_state
+            f.atmt_timeout = timeout
+            f.atmt_timeout._func = f
+            f.atmt_condname = f.__name__
+            return f
+        return deco
+
+    @staticmethod
+    def timer(state, timeout, prio=0):
+        # type: (_StateWrapper, Union[float, int], int) -> Callable[[_StateWrapper, _StateWrapper, Timer], _StateWrapper]  # noqa: E501
+        def deco(f, state=state, timeout=Timer(timeout, prio=prio, autoreload=True)):  # noqa: E501
+            # type: (_StateWrapper, _StateWrapper, Timer) -> _StateWrapper
             f.atmt_type = ATMT.TIMEOUT
             f.atmt_state = state.atmt_state
             f.atmt_timeout = timeout
+            f.atmt_timeout._func = f
             f.atmt_condname = f.__name__
             return f
         return deco
@@ -542,7 +674,7 @@ def __new__(cls, name, bases, dct):
         cls.recv_conditions = {}    # type: Dict[str, List[_StateWrapper]]
         cls.conditions = {}         # type: Dict[str, List[_StateWrapper]]
         cls.ioevents = {}           # type: Dict[str, List[_StateWrapper]]
-        cls.timeout = {}            # type: Dict[str, List[Tuple[int, _StateWrapper]]] # noqa: E501
+        cls.timeout = {}            # type: Dict[str, _TimerList]
         cls.actions = {}            # type: Dict[str, List[_StateWrapper]]
         cls.initial_states = []     # type: List[_StateWrapper]
         cls.stop_states = []        # type: List[_StateWrapper]
@@ -568,7 +700,7 @@ def __new__(cls, name, bases, dct):
                 cls.recv_conditions[s] = []
                 cls.ioevents[s] = []
                 cls.conditions[s] = []
-                cls.timeout[s] = []
+                cls.timeout[s] = _TimerList()
                 if m.atmt_initial:
                     cls.initial_states.append(m)
                 if m.atmt_stop:
@@ -587,14 +719,11 @@ def __new__(cls, name, bases, dct):
                 if m.atmt_as_supersocket is not None:
                     cls.iosupersockets.append(m)
             elif m.atmt_type == ATMT.TIMEOUT:
-                cls.timeout[m.atmt_state].append((m.atmt_timeout, m))
+                cls.timeout[m.atmt_state].add_timer(m.atmt_timeout)
             elif m.atmt_type == ATMT.ACTION:
                 for co in m.atmt_cond:
                     cls.actions[co].append(m)
 
-        for v in six.itervalues(cls.timeout):
-            v.sort(key=lambda x: x[0])
-            v.append((None, None))
         for v in itertools.chain(six.itervalues(cls.conditions),
                                  six.itervalues(cls.recv_conditions),
                                  six.itervalues(cls.ioevents)):
@@ -649,14 +778,14 @@ def build_graph(self):
                         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
-        for k, v2 in six.iteritems(self.timeout):
-            for t, f in v2:
-                if f is None:
-                    continue
-                for n in f.__code__.co_names + f.__code__.co_consts:
+        for k, timers in six.iteritems(self.timeout):
+            for timer in timers:
+                for n in (timer._func.__code__.co_names +
+                          timer._func.__code__.co_consts):
                     if n in self.states:
-                        line = "%s/%.1fs" % (f.atmt_condname, t)
-                        for x in self.actions[f.atmt_condname]:
+                        line = "%s/%.1fs" % (timer._func.atmt_condname,
+                                             timer.get())
+                        for x in self.actions[timer._func.atmt_condname]:
                             line += "\\l>[%s]" % x.__name__
                         s += '\t"%s" -> "%s" [label="%s",color=blue];\n' % (k, n, line)  # noqa: E501
         s += "}\n"
@@ -675,7 +804,7 @@ class Automaton:
     recv_conditions = {}    # type: Dict[str, List[_StateWrapper]]
     conditions = {}         # type: Dict[str, List[_StateWrapper]]
     ioevents = {}           # type: Dict[str, List[_StateWrapper]]
-    timeout = {}            # type: Dict[str, List[Tuple[int, _StateWrapper]]] # noqa: E501
+    timeout = {}            # type: Dict[str, _TimerList]
     actions = {}            # type: Dict[str, List[_StateWrapper]]
     initial_states = []     # type: List[_StateWrapper]
     stop_states = []        # type: List[_StateWrapper]
@@ -748,6 +877,14 @@ def my_send(self, pkt):
         # type: (Packet) -> None
         self.send_sock.send(pkt)
 
+    def timer_by_name(self, name):
+        # type: (str) -> Optional[Timer]
+        for _, timers in six.iteritems(self.timeout):
+            for timer in timers:  # type: Timer
+                if timer._func.atmt_condname == name:
+                    return timer
+        return None
+
     # Utility classes and exceptions
     class _IO_fdwrapper:
         def __init__(self,
@@ -1031,6 +1168,7 @@ def _do_iter(self):
                 elif not isinstance(state_output, list):
                     state_output = state_output,
 
+                timers = self.timeout[self.state.state]
                 # If there are commandMessage, we should skip immediate
                 # conditions.
                 if not select_objects([self.cmdin], 0):
@@ -1041,15 +1179,14 @@ def _do_iter(self):
                     # If still there and no conditions left, we are stuck!
                     if (len(self.recv_conditions[self.state.state]) == 0 and
                         len(self.ioevents[self.state.state]) == 0 and
-                            len(self.timeout[self.state.state]) == 1):
+                            timers.count() == 0):
                         raise self.Stuck("stuck in [%s]" % self.state.state,
                                          state=self.state.state,
                                          result=state_output)
 
                 # Finally listen and pay attention to timeouts
-                expirations = iter(self.timeout[self.state.state])
-                next_timeout, timeout_func = next(expirations)
-                t0 = time.time()
+                timers.reset()
+                time_previous = time.time()
 
                 fds = [self.cmdin]
                 if len(self.recv_conditions[self.state.state]) > 0:
@@ -1057,15 +1194,12 @@ def _do_iter(self):
                 for ioev in self.ioevents[self.state.state]:
                     fds.append(self.ioin[ioev.atmt_ioname])
                 while True:
-                    t = time.time() - t0
-                    if next_timeout is not None:
-                        if next_timeout <= t:
-                            self._run_condition(timeout_func, *state_output)
-                            next_timeout, timeout_func = next(expirations)
-                    if next_timeout is None:
-                        remain = 0
-                    else:
-                        remain = next_timeout - t
+                    time_current = time.time()
+                    timers.decrement(time_current - time_previous)
+                    time_previous = time_current
+                    for timer in timers.expired():
+                        self._run_condition(timer._func, *state_output)
+                    remain = timers.until_next()
 
                     self.debug(5, "Select on %r" % fds)
                     r = select_objects(fds, remain)


Test patch:
diff --git a/test/scapy/automaton.uts b/test/scapy/automaton.uts
--- a/test/scapy/automaton.uts
+++ b/test/scapy/automaton.uts
@@ -352,19 +352,92 @@ r = x
 r
 assert(r == "Venus")
 
+= Automaton timer function
+~ run timers
+
+class TimerTest(Automaton):
+    @ATMT.state(initial=1)
+    def BEGIN(self):
+        self.count1 = 0
+        self.count2 = 0
+    @ATMT.timer(BEGIN, 0.1)
+    def count1(self):
+        self.count1 += 1
+    @ATMT.timer(BEGIN, 0.15)
+    def count2(self):
+        self.count2 += 1
+    @ATMT.timeout(BEGIN, 1)
+    def goto_end(self):
+        raise self.END()
+    @ATMT.state(final=1)
+    def END(self):
+        pass
+
+sm = TimerTest(ll=lambda: None, recvsock=lambda: None)
+sm.run()
+
+assert sm.timer_by_name("count0") is None
+assert sm.timer_by_name("count1") is not None
+assert sm.timer_by_name("count1")._timeout == 0.1
+assert sm.timer_by_name("count2") is not None
+assert sm.timer_by_name("count2")._timeout == 0.15
+assert sm.timer_by_name("goto_end") is not None
+assert sm.timer_by_name("goto_end")._timeout == 1
+assert sm.count1 == 10
+assert sm.count2 == 6
+
+~ reconfigure timers
+
+sm = TimerTest(ll=lambda: None, recvsock=lambda: None)
+sm.timer_by_name("count1").set(0.2)
+sm.timer_by_name("count2").set(0.25)
+sm.run()
+assert sm.count1 == 5
+assert sm.count2 == 4
+
+~ timeout
+
+class TimerTest(Automaton):
+    @ATMT.state(initial=1)
+    def BEGIN(self):
+        self.count1 = 0
+        self.count2 = 0
+    @ATMT.timeout(BEGIN, 0.1)
+    def count1(self):
+        self.count1 += 1
+    @ATMT.timer(BEGIN, 0.15)
+    def count2(self):
+        self.count2 += 1
+    @ATMT.timeout(BEGIN, 1)
+    def goto_end(self):
+        raise self.END()
+    @ATMT.state(final=1)
+    def END(self):
+        pass
+
+sm = TimerTest(ll=lambda: None, recvsock=lambda: None)
+sm.run()
+
+assert sm.count1 == 1
+assert sm.count2 == 6
+
 = Automaton graph
 ~ automaton
 
 class HelloWorld(Automaton):
     @ATMT.state(initial=1)
     def BEGIN(self):
-        pass
-    @ATMT.condition(BEGIN)
-    def wait_for_nothing(self):
+        self.count1 = 0
+        self.count2 = 0
+    @ATMT.timer(BEGIN, 0.1)
+    def count1(self):
+        self.count1 += 1
+    @ATMT.timer(BEGIN, 0.15)
+    def count2(self):
+        self.count2 += 1
+    @ATMT.timeout(BEGIN, 1)
+    def goto_end(self):
         raise self.END()
-    @ATMT.action(wait_for_nothing)
-    def on_nothing(self):
-        pass
     @ATMT.state(final=1)
     def END(self):
         pass