Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/scapy/automaton.py: 36%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# SPDX-License-Identifier: GPL-2.0-only
2# This file is part of Scapy
3# See https://scapy.net/ for more information
4# Copyright (C) Philippe Biondi <phil@secdev.org>
5# Copyright (C) Gabriel Potter
7"""
8Automata with states, transitions and actions.
10TODO:
11 - add documentation for ioevent, as_supersocket...
12"""
14import ctypes
15import itertools
16import logging
17import os
18import random
19import socket
20import sys
21import threading
22import time
23import traceback
24import types
26import select
27from collections import deque
29from scapy.config import conf
30from scapy.consts import WINDOWS
31from scapy.data import MTU
32from scapy.error import log_runtime, warning
33from scapy.interfaces import _GlobInterfaceType
34from scapy.packet import Packet
35from scapy.plist import PacketList
36from scapy.supersocket import SuperSocket, StreamSocket
37from scapy.utils import do_graph
39# Typing imports
40from typing import (
41 Any,
42 Callable,
43 Deque,
44 Dict,
45 Generic,
46 Iterable,
47 Iterator,
48 List,
49 Optional,
50 Set,
51 Tuple,
52 Type,
53 TypeVar,
54 Union,
55 cast,
56)
57from scapy.compat import DecoratorCallable
60# winsock.h
61FD_READ = 0x00000001
64def select_objects(inputs, remain):
65 # type: (Iterable[Any], Union[float, int, None]) -> List[Any]
66 """
67 Select objects. Same than:
68 ``select.select(inputs, [], [], remain)``
70 But also works on Windows, only on objects whose fileno() returns
71 a Windows event. For simplicity, just use `ObjectPipe()` as a queue
72 that you can select on whatever the platform is.
74 If you want an object to be always included in the output of
75 select_objects (i.e. it's not selectable), just make fileno()
76 return a strictly negative value.
78 Example:
80 >>> a, b = ObjectPipe("a"), ObjectPipe("b")
81 >>> b.send("test")
82 >>> select_objects([a, b], 1)
83 [b]
85 :param inputs: objects to process
86 :param remain: timeout. If 0, poll. If None, block.
87 """
88 if not WINDOWS:
89 return select.select(inputs, [], [], remain)[0]
90 inputs = list(inputs)
91 events = []
92 created = []
93 results = set()
94 for i in inputs:
95 if getattr(i, "__selectable_force_select__", False):
96 # Native socket.socket object. We would normally use select.select.
97 evt = ctypes.windll.ws2_32.WSACreateEvent()
98 created.append(evt)
99 res = ctypes.windll.ws2_32.WSAEventSelect(
100 ctypes.c_void_p(i.fileno()),
101 evt,
102 FD_READ
103 )
104 if res == 0:
105 # Was a socket
106 events.append(evt)
107 else:
108 # Fallback to normal event
109 events.append(i.fileno())
110 elif i.fileno() < 0:
111 # Special case: On Windows, we consider that an object that returns
112 # a negative fileno (impossible), is always readable. This is used
113 # in very few places but important (e.g. PcapReader), where we have
114 # no valid fileno (and will stop on EOFError).
115 results.add(i)
116 remain = 0
117 else:
118 events.append(i.fileno())
119 if events:
120 # 0xFFFFFFFF = INFINITE
121 remainms = int(remain * 1000 if remain is not None else 0xFFFFFFFF)
122 if len(events) == 1:
123 res = ctypes.windll.kernel32.WaitForSingleObject(
124 ctypes.c_void_p(events[0]),
125 remainms
126 )
127 else:
128 # Sadly, the only way to emulate select() is to first check
129 # if any object is available using WaitForMultipleObjects
130 # then poll the others.
131 res = ctypes.windll.kernel32.WaitForMultipleObjects(
132 len(events),
133 (ctypes.c_void_p * len(events))(
134 *events
135 ),
136 False,
137 remainms
138 )
139 if res != 0xFFFFFFFF and res != 0x00000102: # Failed or Timeout
140 results.add(inputs[res])
141 if len(events) > 1:
142 # Now poll the others, if any
143 for i, evt in enumerate(events):
144 res = ctypes.windll.kernel32.WaitForSingleObject(
145 ctypes.c_void_p(evt),
146 0 # poll: don't wait
147 )
148 if res == 0:
149 results.add(inputs[i])
150 # Cleanup created events, if any
151 for evt in created:
152 ctypes.windll.ws2_32.WSACloseEvent(evt)
153 return list(results)
156_T = TypeVar("_T")
159class ObjectPipe(Generic[_T]):
160 def __init__(self, name=None):
161 # type: (Optional[str]) -> None
162 self.name = name or "ObjectPipe"
163 self.closed = False
164 self.__rd, self.__wr = os.pipe()
165 self.__queue = deque() # type: Deque[_T]
166 if WINDOWS:
167 self._wincreate()
169 if WINDOWS:
170 def _wincreate(self):
171 # type: () -> None
172 self._fd = cast(int, ctypes.windll.kernel32.CreateEventA(
173 None, True, False,
174 ctypes.create_string_buffer(b"ObjectPipe %f" % random.random())
175 ))
177 def _winset(self):
178 # type: () -> None
179 if ctypes.windll.kernel32.SetEvent(ctypes.c_void_p(self._fd)) == 0:
180 warning(ctypes.FormatError(ctypes.GetLastError()))
182 def _winreset(self):
183 # type: () -> None
184 if ctypes.windll.kernel32.ResetEvent(ctypes.c_void_p(self._fd)) == 0:
185 warning(ctypes.FormatError(ctypes.GetLastError()))
187 def _winclose(self):
188 # type: () -> None
189 if ctypes.windll.kernel32.CloseHandle(ctypes.c_void_p(self._fd)) == 0:
190 warning(ctypes.FormatError(ctypes.GetLastError()))
192 def fileno(self):
193 # type: () -> int
194 if WINDOWS:
195 return self._fd
196 return self.__rd
198 def send(self, obj):
199 # type: (_T) -> int
200 self.__queue.append(obj)
201 if WINDOWS:
202 self._winset()
203 os.write(self.__wr, b"X")
204 return 1
206 def write(self, obj):
207 # type: (_T) -> None
208 self.send(obj)
210 def empty(self):
211 # type: () -> bool
212 return not bool(self.__queue)
214 def flush(self):
215 # type: () -> None
216 pass
218 def recv(self, n=0, options=socket.MsgFlag(0)):
219 # type: (Optional[int], socket.MsgFlag) -> Optional[_T]
220 if self.closed:
221 raise EOFError
222 if options & socket.MSG_PEEK:
223 if self.__queue:
224 return self.__queue[0]
225 return None
226 os.read(self.__rd, 1)
227 elt = self.__queue.popleft()
228 if WINDOWS and not self.__queue:
229 self._winreset()
230 return elt
232 def read(self, n=0):
233 # type: (Optional[int]) -> Optional[_T]
234 return self.recv(n)
236 def clear(self):
237 # type: () -> None
238 if not self.closed:
239 while not self.empty():
240 self.recv()
242 def close(self):
243 # type: () -> None
244 if not self.closed:
245 self.closed = True
246 os.close(self.__rd)
247 os.close(self.__wr)
248 if WINDOWS:
249 try:
250 self._winclose()
251 except ImportError:
252 # Python is shutting down
253 pass
255 def __repr__(self):
256 # type: () -> str
257 return "<%s at %s>" % (self.name, id(self))
259 def __del__(self):
260 # type: () -> None
261 self.close()
263 @staticmethod
264 def select(sockets, remain=conf.recv_poll_rate):
265 # type: (List[SuperSocket], Optional[float]) -> List[SuperSocket]
266 # Only handle ObjectPipes
267 results = []
268 for s in sockets:
269 if s.closed: # allow read to trigger EOF
270 results.append(s)
271 if results:
272 return results
273 return select_objects(sockets, remain)
276class Message:
277 type = None # type: str
278 pkt = None # type: Packet
279 result = None # type: str
280 state = None # type: Message
281 exc_info = None # type: Union[Tuple[None, None, None], Tuple[BaseException, Exception, types.TracebackType]] # noqa: E501
283 def __init__(self, **args):
284 # type: (Any) -> None
285 self.__dict__.update(args)
287 def __repr__(self):
288 # type: () -> str
289 return "<Message %s>" % " ".join(
290 "%s=%r" % (k, v)
291 for k, v in self.__dict__.items()
292 if not k.startswith("_")
293 )
296class Timer():
297 def __init__(self, time, prio=0, autoreload=False):
298 # type: (Union[int, float], int, bool) -> None
299 self._timeout = float(time) # type: float
300 self._time = 0 # type: float
301 self._just_expired = True
302 self._expired = True
303 self._prio = prio
304 self._func = _StateWrapper()
305 self._autoreload = autoreload
307 def get(self):
308 # type: () -> float
309 return self._timeout
311 def set(self, val):
312 # type: (float) -> None
313 self._timeout = val
315 def _reset(self):
316 # type: () -> None
317 self._time = self._timeout
318 self._expired = False
319 self._just_expired = False
321 def _reset_just_expired(self):
322 # type: () -> None
323 self._just_expired = False
325 def _running(self):
326 # type: () -> bool
327 return self._time > 0
329 def _remaining(self):
330 # type: () -> float
331 return max(self._time, 0)
333 def _decrement(self, time):
334 # type: (float) -> None
335 self._time -= time
336 if self._time <= 0:
337 if not self._expired:
338 self._just_expired = True
339 if self._autoreload:
340 # take overshoot into account
341 self._time = self._timeout + self._time
342 else:
343 self._expired = True
344 self._time = 0
346 def __lt__(self, obj):
347 # type: (Timer) -> bool
348 return ((self._time < obj._time) if self._time != obj._time
349 else (self._prio < obj._prio))
351 def __gt__(self, obj):
352 # type: (Timer) -> bool
353 return ((self._time > obj._time) if self._time != obj._time
354 else (self._prio > obj._prio))
356 def __eq__(self, obj):
357 # type: (Any) -> bool
358 if not isinstance(obj, Timer):
359 raise NotImplementedError()
360 return (self._time == obj._time) and (self._prio == obj._prio)
362 def __repr__(self):
363 # type: () -> str
364 return "<Timer %f(%f)>" % (self._time, self._timeout)
367class _TimerList():
368 def __init__(self):
369 # type: () -> None
370 self.timers = [] # type: list[Timer]
372 def add_timer(self, timer):
373 # type: (Timer) -> None
374 self.timers.append(timer)
376 def reset(self):
377 # type: () -> None
378 for t in self.timers:
379 t._reset()
381 def decrement(self, time):
382 # type: (float) -> None
383 for t in self.timers:
384 t._decrement(time)
386 def expired(self):
387 # type: () -> list[Timer]
388 lst = [t for t in self.timers if t._just_expired]
389 lst.sort(key=lambda x: x._prio, reverse=True)
390 for t in lst:
391 t._reset_just_expired()
392 return lst
394 def until_next(self):
395 # type: () -> Optional[float]
396 try:
397 return min([t._remaining() for t in self.timers if t._running()])
398 except ValueError:
399 return None # None means blocking
401 def count(self):
402 # type: () -> int
403 return len(self.timers)
405 def __iter__(self):
406 # type: () -> Iterator[Timer]
407 return self.timers.__iter__()
409 def __repr__(self):
410 # type: () -> str
411 return self.timers.__repr__()
414class _instance_state:
415 def __init__(self, instance):
416 # type: (Any) -> None
417 self.__self__ = instance.__self__
418 self.__func__ = instance.__func__
419 self.__self__.__class__ = instance.__self__.__class__
421 def __getattr__(self, attr):
422 # type: (str) -> Any
423 return getattr(self.__func__, attr)
425 def __call__(self, *args, **kargs):
426 # type: (Any, Any) -> Any
427 return self.__func__(self.__self__, *args, **kargs)
429 def breaks(self):
430 # type: () -> Any
431 return self.__self__.add_breakpoints(self.__func__)
433 def intercepts(self):
434 # type: () -> Any
435 return self.__self__.add_interception_points(self.__func__)
437 def unbreaks(self):
438 # type: () -> Any
439 return self.__self__.remove_breakpoints(self.__func__)
441 def unintercepts(self):
442 # type: () -> Any
443 return self.__self__.remove_interception_points(self.__func__)
446##############
447# Automata #
448##############
450class _StateWrapper:
451 __name__ = None # type: str
452 atmt_type = None # type: str
453 atmt_state = None # type: str
454 atmt_initial = None # type: int
455 atmt_final = None # type: int
456 atmt_stop = None # type: int
457 atmt_error = None # type: int
458 atmt_origfunc = None # type: _StateWrapper
459 atmt_prio = None # type: int
460 atmt_as_supersocket = None # type: Optional[str]
461 atmt_condname = None # type: str
462 atmt_ioname = None # type: str
463 atmt_timeout = None # type: Timer
464 atmt_cond = None # type: Dict[str, int]
465 __code__ = None # type: types.CodeType
466 __call__ = None # type: Callable[..., ATMT.NewStateRequested]
469class ATMT:
470 STATE = "State"
471 ACTION = "Action"
472 CONDITION = "Condition"
473 RECV = "Receive condition"
474 TIMEOUT = "Timeout condition"
475 EOF = "EOF condition"
476 IOEVENT = "I/O event"
478 class NewStateRequested(Exception):
479 def __init__(self, state_func, automaton, *args, **kargs):
480 # type: (Any, ATMT, Any, Any) -> None
481 self.func = state_func
482 self.state = state_func.atmt_state
483 self.initial = state_func.atmt_initial
484 self.error = state_func.atmt_error
485 self.stop = state_func.atmt_stop
486 self.final = state_func.atmt_final
487 Exception.__init__(self, "Request state [%s]" % self.state)
488 self.automaton = automaton
489 self.args = args
490 self.kargs = kargs
491 self.action_parameters() # init action parameters
493 def action_parameters(self, *args, **kargs):
494 # type: (Any, Any) -> ATMT.NewStateRequested
495 self.action_args = args
496 self.action_kargs = kargs
497 return self
499 def run(self):
500 # type: () -> Any
501 return self.func(self.automaton, *self.args, **self.kargs)
503 def __repr__(self):
504 # type: () -> str
505 return "NewStateRequested(%s)" % self.state
507 @staticmethod
508 def state(initial=0, # type: int
509 final=0, # type: int
510 stop=0, # type: int
511 error=0 # type: int
512 ):
513 # type: (...) -> Callable[[DecoratorCallable], DecoratorCallable]
514 def deco(f, initial=initial, final=final):
515 # type: (_StateWrapper, int, int) -> _StateWrapper
516 f.atmt_type = ATMT.STATE
517 f.atmt_state = f.__name__
518 f.atmt_initial = initial
519 f.atmt_final = final
520 f.atmt_stop = stop
521 f.atmt_error = error
523 def _state_wrapper(self, *args, **kargs):
524 # type: (ATMT, Any, Any) -> ATMT.NewStateRequested
525 return ATMT.NewStateRequested(f, self, *args, **kargs)
527 state_wrapper = cast(_StateWrapper, _state_wrapper)
528 state_wrapper.__name__ = "%s_wrapper" % f.__name__
529 state_wrapper.atmt_type = ATMT.STATE
530 state_wrapper.atmt_state = f.__name__
531 state_wrapper.atmt_initial = initial
532 state_wrapper.atmt_final = final
533 state_wrapper.atmt_stop = stop
534 state_wrapper.atmt_error = error
535 state_wrapper.atmt_origfunc = f
536 return state_wrapper
537 return deco # type: ignore
539 @staticmethod
540 def action(cond, prio=0):
541 # type: (Any, int) -> Callable[[_StateWrapper, _StateWrapper], _StateWrapper] # noqa: E501
542 def deco(f, cond=cond):
543 # type: (_StateWrapper, _StateWrapper) -> _StateWrapper
544 if not hasattr(f, "atmt_type"):
545 f.atmt_cond = {}
546 f.atmt_type = ATMT.ACTION
547 f.atmt_cond[cond.atmt_condname] = prio
548 return f
549 return deco
551 @staticmethod
552 def condition(state, prio=0):
553 # type: (Any, int) -> Callable[[_StateWrapper, _StateWrapper], _StateWrapper] # noqa: E501
554 def deco(f, state=state):
555 # type: (_StateWrapper, _StateWrapper) -> Any
556 f.atmt_type = ATMT.CONDITION
557 f.atmt_state = state.atmt_state
558 f.atmt_condname = f.__name__
559 f.atmt_prio = prio
560 return f
561 return deco
563 @staticmethod
564 def receive_condition(state, prio=0):
565 # type: (_StateWrapper, int) -> Callable[[_StateWrapper, _StateWrapper], _StateWrapper] # noqa: E501
566 def deco(f, state=state):
567 # type: (_StateWrapper, _StateWrapper) -> _StateWrapper
568 f.atmt_type = ATMT.RECV
569 f.atmt_state = state.atmt_state
570 f.atmt_condname = f.__name__
571 f.atmt_prio = prio
572 return f
573 return deco
575 @staticmethod
576 def ioevent(state, # type: _StateWrapper
577 name, # type: str
578 prio=0, # type: int
579 as_supersocket=None # type: Optional[str]
580 ):
581 # type: (...) -> Callable[[_StateWrapper, _StateWrapper], _StateWrapper] # noqa: E501
582 def deco(f, state=state):
583 # type: (_StateWrapper, _StateWrapper) -> _StateWrapper
584 f.atmt_type = ATMT.IOEVENT
585 f.atmt_state = state.atmt_state
586 f.atmt_condname = f.__name__
587 f.atmt_ioname = name
588 f.atmt_prio = prio
589 f.atmt_as_supersocket = as_supersocket
590 return f
591 return deco
593 @staticmethod
594 def timeout(state, timeout):
595 # type: (_StateWrapper, Union[int, float]) -> Callable[[_StateWrapper, _StateWrapper, Timer], _StateWrapper] # noqa: E501
596 def deco(f, state=state, timeout=Timer(timeout)):
597 # type: (_StateWrapper, _StateWrapper, Timer) -> _StateWrapper
598 f.atmt_type = ATMT.TIMEOUT
599 f.atmt_state = state.atmt_state
600 f.atmt_timeout = timeout
601 f.atmt_timeout._func = f
602 f.atmt_condname = f.__name__
603 return f
604 return deco
606 @staticmethod
607 def timer(state, timeout, prio=0):
608 # type: (_StateWrapper, Union[float, int], int) -> Callable[[_StateWrapper, _StateWrapper, Timer], _StateWrapper] # noqa: E501
609 def deco(f, state=state, timeout=Timer(timeout, prio=prio, autoreload=True)):
610 # type: (_StateWrapper, _StateWrapper, Timer) -> _StateWrapper
611 f.atmt_type = ATMT.TIMEOUT
612 f.atmt_state = state.atmt_state
613 f.atmt_timeout = timeout
614 f.atmt_timeout._func = f
615 f.atmt_condname = f.__name__
616 return f
617 return deco
619 @staticmethod
620 def eof(state):
621 # type: (_StateWrapper) -> Callable[[_StateWrapper, _StateWrapper], _StateWrapper] # noqa: E501
622 def deco(f, state=state):
623 # type: (_StateWrapper, _StateWrapper) -> _StateWrapper
624 f.atmt_type = ATMT.EOF
625 f.atmt_state = state.atmt_state
626 f.atmt_condname = f.__name__
627 return f
628 return deco
631class _ATMT_Command:
632 RUN = "RUN"
633 NEXT = "NEXT"
634 FREEZE = "FREEZE"
635 STOP = "STOP"
636 FORCESTOP = "FORCESTOP"
637 END = "END"
638 EXCEPTION = "EXCEPTION"
639 SINGLESTEP = "SINGLESTEP"
640 BREAKPOINT = "BREAKPOINT"
641 INTERCEPT = "INTERCEPT"
642 ACCEPT = "ACCEPT"
643 REPLACE = "REPLACE"
644 REJECT = "REJECT"
647class _ATMT_supersocket(SuperSocket):
648 def __init__(self,
649 name, # type: str
650 ioevent, # type: str
651 automaton, # type: Type[Automaton]
652 proto, # type: Callable[[bytes], Any]
653 *args, # type: Any
654 **kargs # type: Any
655 ):
656 # type: (...) -> None
657 self.name = name
658 self.ioevent = ioevent
659 self.proto = proto
660 # write, read
661 self.spa, self.spb = ObjectPipe[Any]("spa"), \
662 ObjectPipe[Any]("spb")
663 kargs["external_fd"] = {ioevent: (self.spa, self.spb)}
664 kargs["is_atmt_socket"] = True
665 kargs["atmt_socket"] = self.name
666 self.atmt = automaton(*args, **kargs)
667 self.atmt.runbg()
669 def send(self, s):
670 # type: (Any) -> int
671 return self.spa.send(s)
673 def fileno(self):
674 # type: () -> int
675 return self.spb.fileno()
677 # note: _ATMT_supersocket may return bytes in certain cases, which
678 # is expected. We cheat on typing.
679 def recv(self, n=MTU, **kwargs): # type: ignore
680 # type: (int, **Any) -> Any
681 r = self.spb.recv(n)
682 if self.proto is not None and r is not None:
683 r = self.proto(r, **kwargs)
684 if self.atmt.atmt_session is not None:
685 # Apply session if provided
686 r = self.atmt.atmt_session.process(r)
687 return r
689 def close(self):
690 # type: () -> None
691 if not self.closed:
692 self.atmt.stop()
693 self.atmt.destroy()
694 self.spa.close()
695 self.spb.close()
696 self.closed = True
698 @staticmethod
699 def select(sockets, remain=conf.recv_poll_rate):
700 # type: (List[SuperSocket], Optional[float]) -> List[SuperSocket]
701 return select_objects(sockets, remain)
704class _ATMT_to_supersocket:
705 def __init__(self, name, ioevent, automaton):
706 # type: (str, str, Type[Automaton]) -> None
707 self.name = name
708 self.ioevent = ioevent
709 self.automaton = automaton
711 def __call__(self, proto, *args, **kargs):
712 # type: (Callable[[bytes], Any], Any, Any) -> _ATMT_supersocket
713 return _ATMT_supersocket(
714 self.name, self.ioevent, self.automaton,
715 proto, *args, **kargs
716 )
719class Automaton_metaclass(type):
720 def __new__(cls, name, bases, dct):
721 # type: (str, Tuple[Any], Dict[str, Any]) -> Type[Automaton]
722 cls = super(Automaton_metaclass, cls).__new__( # type: ignore
723 cls, name, bases, dct
724 )
725 cls.states = {}
726 cls.recv_conditions = {} # type: Dict[str, List[_StateWrapper]]
727 cls.conditions = {} # type: Dict[str, List[_StateWrapper]]
728 cls.ioevents = {} # type: Dict[str, List[_StateWrapper]]
729 cls.timeout = {} # type: Dict[str, _TimerList]
730 cls.eofs = {} # type: Dict[str, _StateWrapper]
731 cls.actions = {} # type: Dict[str, List[_StateWrapper]]
732 cls.initial_states = [] # type: List[_StateWrapper]
733 cls.stop_state = None # type: Optional[_StateWrapper]
734 cls.ionames = []
735 cls.iosupersockets = []
737 members = {}
738 classes = [cls]
739 while classes:
740 c = classes.pop(0) # order is important to avoid breaking method overloading # noqa: E501
741 classes += list(c.__bases__)
742 for k, v in c.__dict__.items(): # type: ignore
743 if k not in members:
744 members[k] = v
746 decorated = [v for v in members.values()
747 if hasattr(v, "atmt_type")]
749 for m in decorated:
750 if m.atmt_type == ATMT.STATE:
751 s = m.atmt_state
752 cls.states[s] = m
753 cls.recv_conditions[s] = []
754 cls.ioevents[s] = []
755 cls.conditions[s] = []
756 cls.timeout[s] = _TimerList()
757 if m.atmt_initial:
758 cls.initial_states.append(m)
759 if m.atmt_stop:
760 if cls.stop_state is not None:
761 raise ValueError("There can only be a single stop state !")
762 cls.stop_state = m
763 elif m.atmt_type in [ATMT.CONDITION, ATMT.RECV, ATMT.TIMEOUT, ATMT.IOEVENT, ATMT.EOF]: # noqa: E501
764 cls.actions[m.atmt_condname] = []
766 for m in decorated:
767 if m.atmt_type == ATMT.CONDITION:
768 cls.conditions[m.atmt_state].append(m)
769 elif m.atmt_type == ATMT.RECV:
770 cls.recv_conditions[m.atmt_state].append(m)
771 elif m.atmt_type == ATMT.EOF:
772 cls.eofs[m.atmt_state] = m
773 elif m.atmt_type == ATMT.IOEVENT:
774 cls.ioevents[m.atmt_state].append(m)
775 cls.ionames.append(m.atmt_ioname)
776 if m.atmt_as_supersocket is not None:
777 cls.iosupersockets.append(m)
778 elif m.atmt_type == ATMT.TIMEOUT:
779 cls.timeout[m.atmt_state].add_timer(m.atmt_timeout)
780 elif m.atmt_type == ATMT.ACTION:
781 for co in m.atmt_cond:
782 cls.actions[co].append(m)
784 for v in itertools.chain(
785 cls.conditions.values(),
786 cls.recv_conditions.values(),
787 cls.ioevents.values()
788 ):
789 v.sort(key=lambda x: x.atmt_prio)
790 for condname, actlst in cls.actions.items():
791 actlst.sort(key=lambda x: x.atmt_cond[condname])
793 for ioev in cls.iosupersockets:
794 setattr(cls, ioev.atmt_as_supersocket,
795 _ATMT_to_supersocket(
796 ioev.atmt_as_supersocket,
797 ioev.atmt_ioname,
798 cast(Type["Automaton"], cls)))
800 # Inject signature
801 try:
802 import inspect
803 cls.__signature__ = inspect.signature(cls.parse_args) # type: ignore # noqa: E501
804 except (ImportError, AttributeError):
805 pass
807 return cast(Type["Automaton"], cls)
809 def build_graph(self):
810 # type: () -> str
811 s = 'digraph "%s" {\n' % self.__class__.__name__
813 se = "" # Keep initial nodes at the beginning for better rendering
814 for st in self.states.values():
815 if st.atmt_initial:
816 se = ('\t"%s" [ style=filled, fillcolor=blue, shape=box, root=true];\n' % st.atmt_state) + se # noqa: E501
817 elif st.atmt_final:
818 se += '\t"%s" [ style=filled, fillcolor=green, shape=octagon ];\n' % st.atmt_state # noqa: E501
819 elif st.atmt_error:
820 se += '\t"%s" [ style=filled, fillcolor=red, shape=octagon ];\n' % st.atmt_state # noqa: E501
821 elif st.atmt_stop:
822 se += '\t"%s" [ style=filled, fillcolor=orange, shape=box, root=true ];\n' % st.atmt_state # noqa: E501
823 s += se
825 for st in self.states.values():
826 names = list(
827 st.atmt_origfunc.__code__.co_names +
828 st.atmt_origfunc.__code__.co_consts
829 )
830 while names:
831 n = names.pop()
832 if n in self.states:
833 s += '\t"%s" -> "%s" [ color=green ];\n' % (st.atmt_state, n)
834 elif n in self.__dict__:
835 # function indirection
836 if callable(self.__dict__[n]):
837 names.extend(self.__dict__[n].__code__.co_names)
838 names.extend(self.__dict__[n].__code__.co_consts)
840 for c, sty, k, v in (
841 [("purple", "solid", k, v) for k, v in self.conditions.items()] +
842 [("red", "solid", k, v) for k, v in self.recv_conditions.items()] +
843 [("orange", "solid", k, v) for k, v in self.ioevents.items()] +
844 [("black", "dashed", k, [v]) for k, v in self.eofs.items()]
845 ):
846 for f in v:
847 names = list(f.__code__.co_names + f.__code__.co_consts)
848 while names:
849 n = names.pop()
850 if n in self.states:
851 line = f.atmt_condname
852 for x in self.actions[f.atmt_condname]:
853 line += "\\l>[%s]" % x.__name__
854 s += '\t"%s" -> "%s" [label="%s", color=%s, style=%s];\n' % (
855 k,
856 n,
857 line,
858 c,
859 sty,
860 )
861 elif n in self.__dict__:
862 # function indirection
863 if callable(self.__dict__[n]) and hasattr(self.__dict__[n], "__code__"): # noqa: E501
864 names.extend(self.__dict__[n].__code__.co_names)
865 names.extend(self.__dict__[n].__code__.co_consts)
866 for k, timers in self.timeout.items():
867 for timer in timers:
868 for n in (timer._func.__code__.co_names +
869 timer._func.__code__.co_consts):
870 if n in self.states:
871 line = "%s/%.1fs" % (timer._func.atmt_condname,
872 timer.get())
873 for x in self.actions[timer._func.atmt_condname]:
874 line += "\\l>[%s]" % x.__name__
875 s += '\t"%s" -> "%s" [label="%s",color=blue];\n' % (k, n, line) # noqa: E501
876 s += "}\n"
877 return s
879 def graph(self, **kargs):
880 # type: (Any) -> Optional[str]
881 s = self.build_graph()
882 return do_graph(s, **kargs)
885class Automaton(metaclass=Automaton_metaclass):
886 states = {} # type: Dict[str, _StateWrapper]
887 state = None # type: ATMT.NewStateRequested
888 recv_conditions = {} # type: Dict[str, List[_StateWrapper]]
889 conditions = {} # type: Dict[str, List[_StateWrapper]]
890 eofs = {} # type: Dict[str, _StateWrapper]
891 ioevents = {} # type: Dict[str, List[_StateWrapper]]
892 timeout = {} # type: Dict[str, _TimerList]
893 actions = {} # type: Dict[str, List[_StateWrapper]]
894 initial_states = [] # type: List[_StateWrapper]
895 stop_state = None # type: Optional[_StateWrapper]
896 ionames = [] # type: List[str]
897 iosupersockets = [] # type: List[SuperSocket]
899 # used for spawn()
900 pkt_cls = conf.raw_layer
901 socketcls = StreamSocket
903 # Internals
904 def __init__(self, *args, **kargs):
905 # type: (Any, Any) -> None
906 external_fd = kargs.pop("external_fd", {})
907 if "sock" in kargs:
908 # We use a bi-directional sock
909 self.sock = kargs["sock"]
910 else:
911 # Separate sockets
912 self.sock = None
913 self.send_sock_class = kargs.pop("ll", conf.L3socket)
914 self.recv_sock_class = kargs.pop("recvsock", conf.L2listen)
915 self.listen_sock = None # type: Optional[SuperSocket]
916 self.send_sock = None # type: Optional[SuperSocket]
917 self.is_atmt_socket = kargs.pop("is_atmt_socket", False)
918 self.atmt_socket = kargs.pop("atmt_socket", None)
919 self.started = threading.Lock()
920 self.threadid = None # type: Optional[int]
921 self.breakpointed = None
922 self.breakpoints = set() # type: Set[_StateWrapper]
923 self.interception_points = set() # type: Set[_StateWrapper]
924 self.intercepted_packet = None # type: Union[None, Packet]
925 self.debug_level = 0
926 self.init_args = args
927 self.init_kargs = kargs
928 self.io = type.__new__(type, "IOnamespace", (), {})
929 self.oi = type.__new__(type, "IOnamespace", (), {})
930 self.cmdin = ObjectPipe[Message]("cmdin")
931 self.cmdout = ObjectPipe[Message]("cmdout")
932 self.ioin = {}
933 self.ioout = {}
934 self.packets = PacketList() # type: PacketList
935 self.atmt_session = kargs.pop("session", None)
936 for n in self.__class__.ionames:
937 extfd = external_fd.get(n)
938 if not isinstance(extfd, tuple):
939 extfd = (extfd, extfd)
940 ioin, ioout = extfd
941 if ioin is None:
942 ioin = ObjectPipe("ioin")
943 else:
944 ioin = self._IO_fdwrapper(ioin, None)
945 if ioout is None:
946 ioout = ObjectPipe("ioout")
947 else:
948 ioout = self._IO_fdwrapper(None, ioout)
950 self.ioin[n] = ioin
951 self.ioout[n] = ioout
952 ioin.ioname = n
953 ioout.ioname = n
954 setattr(self.io, n, self._IO_mixer(ioout, ioin))
955 setattr(self.oi, n, self._IO_mixer(ioin, ioout))
957 for stname in self.states:
958 setattr(self, stname,
959 _instance_state(getattr(self, stname)))
961 self.start()
963 def parse_args(self, debug=0, store=0, session=None, **kargs):
964 # type: (int, int, Any, Any) -> None
965 self.debug_level = debug
966 if debug:
967 conf.logLevel = logging.DEBUG
968 if session:
969 self.atmt_session = session
970 self.socket_kargs = kargs
971 self.store_packets = store
973 @classmethod
974 def spawn(cls,
975 port: int,
976 iface: Optional[_GlobInterfaceType] = None,
977 local_ip: Optional[str] = None,
978 bg: bool = False,
979 **kwargs: Any) -> Optional[socket.socket]:
980 """
981 Spawn a TCP server that listens for connections and start the automaton
982 for each new client.
984 :param port: the port to listen to
985 :param bg: background mode? (default: False)
987 Note that in background mode, you shall close the TCP server as such::
989 srv = MyAutomaton.spawn(8080, bg=True)
990 srv.shutdown(socket.SHUT_RDWR) # important
991 srv.close()
992 """
993 from scapy.arch import get_if_addr
994 # create server sock and bind it
995 ssock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
996 if local_ip is None:
997 local_ip = get_if_addr(iface or conf.iface)
998 try:
999 ssock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
1000 except OSError:
1001 pass
1002 ssock.bind((local_ip, port))
1003 ssock.listen(5)
1004 clients = []
1005 if kwargs.get("verb", True):
1006 print(conf.color_theme.green(
1007 "Server %s started listening on %s" % (
1008 cls.__name__,
1009 (local_ip, port),
1010 )
1011 ))
1013 def _run() -> None:
1014 # Wait for clients forever
1015 try:
1016 while True:
1017 atmt_server = None
1018 clientsocket, address = ssock.accept()
1019 if kwargs.get("verb", True):
1020 print(conf.color_theme.gold(
1021 "\u2503 Connection received from %s" % repr(address)
1022 ))
1023 try:
1024 # Start atmt class with socket
1025 if cls.socketcls is not None:
1026 sock = cls.socketcls(clientsocket, cls.pkt_cls)
1027 else:
1028 sock = clientsocket
1029 atmt_server = cls(
1030 sock=sock,
1031 iface=iface, **kwargs
1032 )
1033 except OSError:
1034 if atmt_server is not None:
1035 atmt_server.destroy()
1036 if kwargs.get("verb", True):
1037 print("X Connection aborted.")
1038 if kwargs.get("debug", 0) > 0:
1039 traceback.print_exc()
1040 continue
1041 clients.append((atmt_server, clientsocket))
1042 # start atmt
1043 atmt_server.runbg()
1044 # housekeeping
1045 for atmt, clientsocket in clients:
1046 if not atmt.isrunning():
1047 atmt.destroy()
1048 except KeyboardInterrupt:
1049 print("X Exiting.")
1050 ssock.shutdown(socket.SHUT_RDWR)
1051 except OSError:
1052 print("X Server closed.")
1053 if kwargs.get("debug", 0) > 0:
1054 traceback.print_exc()
1055 finally:
1056 for atmt, clientsocket in clients:
1057 try:
1058 atmt.forcestop(wait=False)
1059 atmt.destroy()
1060 except Exception:
1061 pass
1062 try:
1063 clientsocket.shutdown(socket.SHUT_RDWR)
1064 clientsocket.close()
1065 except Exception:
1066 pass
1067 ssock.close()
1068 if bg:
1069 # Background
1070 threading.Thread(target=_run).start()
1071 return ssock
1072 else:
1073 # Non-background
1074 _run()
1075 return None
1077 def master_filter(self, pkt):
1078 # type: (Packet) -> bool
1079 return True
1081 def my_send(self, pkt, **kwargs):
1082 # type: (Packet, **Any) -> None
1083 if not self.send_sock:
1084 raise ValueError("send_sock is None !")
1085 self.send_sock.send(pkt, **kwargs)
1087 def update_sock(self, sock):
1088 # type: (SuperSocket) -> None
1089 """
1090 Update the socket used by the automata.
1091 Typically used in an eof event to reconnect.
1092 """
1093 self.sock = sock
1094 if self.listen_sock is not None:
1095 self.listen_sock = self.sock
1096 if self.send_sock:
1097 self.send_sock = self.sock
1099 def timer_by_name(self, name):
1100 # type: (str) -> Optional[Timer]
1101 for _, timers in self.timeout.items():
1102 for timer in timers: # type: Timer
1103 if timer._func.atmt_condname == name:
1104 return timer
1105 return None
1107 # Utility classes and exceptions
1108 class _IO_fdwrapper:
1109 def __init__(self,
1110 rd, # type: Union[int, ObjectPipe[bytes], None]
1111 wr # type: Union[int, ObjectPipe[bytes], None]
1112 ):
1113 # type: (...) -> None
1114 self.rd = rd
1115 self.wr = wr
1116 if isinstance(self.rd, socket.socket):
1117 self.__selectable_force_select__ = True
1119 def fileno(self):
1120 # type: () -> int
1121 if isinstance(self.rd, int):
1122 return self.rd
1123 elif self.rd:
1124 return self.rd.fileno()
1125 return 0
1127 def read(self, n=65535):
1128 # type: (int) -> Optional[bytes]
1129 if isinstance(self.rd, int):
1130 return os.read(self.rd, n)
1131 elif self.rd:
1132 return self.rd.recv(n)
1133 return None
1135 def write(self, msg):
1136 # type: (bytes) -> int
1137 if isinstance(self.wr, int):
1138 return os.write(self.wr, msg)
1139 elif self.wr:
1140 return self.wr.send(msg)
1141 return 0
1143 def recv(self, n=65535):
1144 # type: (int) -> Optional[bytes]
1145 return self.read(n)
1147 def send(self, msg):
1148 # type: (bytes) -> int
1149 return self.write(msg)
1151 class _IO_mixer:
1152 def __init__(self,
1153 rd, # type: ObjectPipe[Any]
1154 wr, # type: ObjectPipe[Any]
1155 ):
1156 # type: (...) -> None
1157 self.rd = rd
1158 self.wr = wr
1160 def fileno(self):
1161 # type: () -> Any
1162 if isinstance(self.rd, ObjectPipe):
1163 return self.rd.fileno()
1164 return self.rd
1166 def recv(self, n=None):
1167 # type: (Optional[int]) -> Any
1168 return self.rd.recv(n)
1170 def read(self, n=None):
1171 # type: (Optional[int]) -> Any
1172 return self.recv(n)
1174 def send(self, msg):
1175 # type: (str) -> int
1176 return self.wr.send(msg)
1178 def write(self, msg):
1179 # type: (str) -> int
1180 return self.send(msg)
1182 class AutomatonException(Exception):
1183 def __init__(self, msg, state=None, result=None):
1184 # type: (str, Optional[Message], Optional[str]) -> None
1185 Exception.__init__(self, msg)
1186 self.state = state
1187 self.result = result
1189 class AutomatonError(AutomatonException):
1190 pass
1192 class ErrorState(AutomatonException):
1193 pass
1195 class Stuck(AutomatonException):
1196 pass
1198 class AutomatonStopped(AutomatonException):
1199 pass
1201 class Breakpoint(AutomatonStopped):
1202 pass
1204 class Singlestep(AutomatonStopped):
1205 pass
1207 class InterceptionPoint(AutomatonStopped):
1208 def __init__(self, msg, state=None, result=None, packet=None):
1209 # type: (str, Optional[Message], Optional[str], Optional[Packet]) -> None
1210 Automaton.AutomatonStopped.__init__(self, msg, state=state, result=result)
1211 self.packet = packet
1213 class CommandMessage(AutomatonException):
1214 pass
1216 # Services
1217 def debug(self, lvl, msg):
1218 # type: (int, str) -> None
1219 if self.debug_level >= lvl:
1220 log_runtime.debug(msg)
1222 def isrunning(self):
1223 # type: () -> bool
1224 return self.started.locked()
1226 def send(self, pkt, **kwargs):
1227 # type: (Packet, **Any) -> None
1228 if self.state.state in self.interception_points:
1229 self.debug(3, "INTERCEPT: packet intercepted: %s" % pkt.summary())
1230 self.intercepted_packet = pkt
1231 self.cmdout.send(
1232 Message(type=_ATMT_Command.INTERCEPT,
1233 state=self.state, pkt=pkt)
1234 )
1235 cmd = self.cmdin.recv()
1236 if not cmd:
1237 self.debug(3, "CANCELLED")
1238 return
1239 self.intercepted_packet = None
1240 if cmd.type == _ATMT_Command.REJECT:
1241 self.debug(3, "INTERCEPT: packet rejected")
1242 return
1243 elif cmd.type == _ATMT_Command.REPLACE:
1244 pkt = cmd.pkt
1245 self.debug(3, "INTERCEPT: packet replaced by: %s" % pkt.summary()) # noqa: E501
1246 elif cmd.type == _ATMT_Command.ACCEPT:
1247 self.debug(3, "INTERCEPT: packet accepted")
1248 else:
1249 raise self.AutomatonError("INTERCEPT: unknown verdict: %r" % cmd.type) # noqa: E501
1250 self.my_send(pkt, **kwargs)
1251 self.debug(3, "SENT : %s" % pkt.summary())
1253 if self.store_packets:
1254 self.packets.append(pkt.copy())
1256 def __iter__(self):
1257 # type: () -> Automaton
1258 return self
1260 def __del__(self):
1261 # type: () -> None
1262 self.destroy()
1264 def _run_condition(self, cond, *args, **kargs):
1265 # type: (_StateWrapper, Any, Any) -> None
1266 try:
1267 self.debug(5, "Trying %s [%s]" % (cond.atmt_type, cond.atmt_condname)) # noqa: E501
1268 cond(self, *args, **kargs)
1269 except ATMT.NewStateRequested as state_req:
1270 self.debug(2, "%s [%s] taken to state [%s]" % (cond.atmt_type, cond.atmt_condname, state_req.state)) # noqa: E501
1271 if cond.atmt_type == ATMT.RECV:
1272 if self.store_packets:
1273 self.packets.append(args[0])
1274 for action in self.actions[cond.atmt_condname]:
1275 self.debug(2, " + Running action [%s]" % action.__name__)
1276 action(self, *state_req.action_args, **state_req.action_kargs)
1277 raise
1278 except Exception as e:
1279 self.debug(2, "%s [%s] raised exception [%s]" % (cond.atmt_type, cond.atmt_condname, e)) # noqa: E501
1280 raise
1281 else:
1282 self.debug(2, "%s [%s] not taken" % (cond.atmt_type, cond.atmt_condname)) # noqa: E501
1284 def _do_start(self, *args, **kargs):
1285 # type: (Any, Any) -> None
1286 ready = threading.Event()
1287 _t = threading.Thread(
1288 target=self._do_control,
1289 args=(ready,) + (args),
1290 kwargs=kargs,
1291 name="scapy.automaton _do_start"
1292 )
1293 _t.daemon = True
1294 _t.start()
1295 ready.wait()
1297 def _do_control(self, ready, *args, **kargs):
1298 # type: (threading.Event, Any, Any) -> None
1299 with self.started:
1300 self.threadid = threading.current_thread().ident
1301 if self.threadid is None:
1302 self.threadid = 0
1304 # Update default parameters
1305 a = args + self.init_args[len(args):]
1306 k = self.init_kargs.copy()
1307 k.update(kargs)
1308 self.parse_args(*a, **k)
1310 # Start the automaton
1311 self.state = self.initial_states[0](self)
1312 self.send_sock = self.sock or self.send_sock_class(**self.socket_kargs)
1313 if self.recv_conditions:
1314 # Only start a receiving socket if we have at least one recv_conditions
1315 self.listen_sock = self.sock or self.recv_sock_class(**self.socket_kargs) # noqa: E501
1316 self.packets = PacketList(name="session[%s]" % self.__class__.__name__)
1318 singlestep = True
1319 iterator = self._do_iter()
1320 self.debug(3, "Starting control thread [tid=%i]" % self.threadid)
1321 # Sync threads
1322 ready.set()
1323 try:
1324 while True:
1325 c = self.cmdin.recv()
1326 if c is None:
1327 return None
1328 self.debug(5, "Received command %s" % c.type)
1329 if c.type == _ATMT_Command.RUN:
1330 singlestep = False
1331 elif c.type == _ATMT_Command.NEXT:
1332 singlestep = True
1333 elif c.type == _ATMT_Command.FREEZE:
1334 continue
1335 elif c.type == _ATMT_Command.STOP:
1336 if self.stop_state:
1337 # There is a stop state
1338 self.state = self.stop_state()
1339 iterator = self._do_iter()
1340 else:
1341 # Act as FORCESTOP
1342 break
1343 elif c.type == _ATMT_Command.FORCESTOP:
1344 break
1345 while True:
1346 state = next(iterator)
1347 if isinstance(state, self.CommandMessage):
1348 break
1349 elif isinstance(state, self.Breakpoint):
1350 c = Message(type=_ATMT_Command.BREAKPOINT, state=state) # noqa: E501
1351 self.cmdout.send(c)
1352 break
1353 if singlestep:
1354 c = Message(type=_ATMT_Command.SINGLESTEP, state=state) # noqa: E501
1355 self.cmdout.send(c)
1356 break
1357 except (StopIteration, RuntimeError):
1358 c = Message(type=_ATMT_Command.END,
1359 result=self.final_state_output)
1360 self.cmdout.send(c)
1361 except Exception as e:
1362 exc_info = sys.exc_info()
1363 self.debug(3, "Transferring exception from tid=%i:\n%s" % (self.threadid, "".join(traceback.format_exception(*exc_info)))) # noqa: E501
1364 m = Message(type=_ATMT_Command.EXCEPTION, exception=e, exc_info=exc_info) # noqa: E501
1365 self.cmdout.send(m)
1366 self.debug(3, "Stopping control thread (tid=%i)" % self.threadid)
1367 self.threadid = None
1368 if self.listen_sock:
1369 self.listen_sock.close()
1370 if self.send_sock:
1371 self.send_sock.close()
1373 def _do_iter(self):
1374 # type: () -> Iterator[Union[Automaton.AutomatonException, Automaton.AutomatonStopped, ATMT.NewStateRequested, None]] # noqa: E501
1375 while True:
1376 try:
1377 self.debug(1, "## state=[%s]" % self.state.state)
1379 # Entering a new state. First, call new state function
1380 if self.state.state in self.breakpoints and self.state.state != self.breakpointed: # noqa: E501
1381 self.breakpointed = self.state.state
1382 yield self.Breakpoint("breakpoint triggered on state %s" % self.state.state, # noqa: E501
1383 state=self.state.state)
1384 self.breakpointed = None
1385 state_output = self.state.run()
1386 if self.state.error:
1387 raise self.ErrorState("Reached %s: [%r]" % (self.state.state, state_output), # noqa: E501
1388 result=state_output, state=self.state.state) # noqa: E501
1389 if self.state.final:
1390 self.final_state_output = state_output
1391 return
1393 if state_output is None:
1394 state_output = ()
1395 elif not isinstance(state_output, list):
1396 state_output = state_output,
1398 timers = self.timeout[self.state.state]
1399 # If there are commandMessage, we should skip immediate
1400 # conditions.
1401 if not select_objects([self.cmdin], 0):
1402 # Then check immediate conditions
1403 for cond in self.conditions[self.state.state]:
1404 self._run_condition(cond, *state_output)
1406 # If still there and no conditions left, we are stuck!
1407 if (len(self.recv_conditions[self.state.state]) == 0 and
1408 len(self.ioevents[self.state.state]) == 0 and
1409 timers.count() == 0):
1410 raise self.Stuck("stuck in [%s]" % self.state.state,
1411 state=self.state.state,
1412 result=state_output)
1414 # Finally listen and pay attention to timeouts
1415 timers.reset()
1416 time_previous = time.time()
1418 fds = [self.cmdin] # type: List[Union[SuperSocket, ObjectPipe[Any]]]
1419 select_func = select_objects
1420 if self.listen_sock and self.recv_conditions[self.state.state]:
1421 fds.append(self.listen_sock)
1422 select_func = self.listen_sock.select # type: ignore
1423 for ioev in self.ioevents[self.state.state]:
1424 fds.append(self.ioin[ioev.atmt_ioname])
1425 while True:
1426 time_current = time.time()
1427 timers.decrement(time_current - time_previous)
1428 time_previous = time_current
1429 for timer in timers.expired():
1430 self._run_condition(timer._func, *state_output)
1431 remain = timers.until_next()
1433 self.debug(5, "Select on %r" % fds)
1434 r = select_func(fds, remain)
1435 self.debug(5, "Selected %r" % r)
1436 for fd in r:
1437 self.debug(5, "Looking at %r" % fd)
1438 if fd == self.cmdin:
1439 yield self.CommandMessage("Received command message") # noqa: E501
1440 elif fd == self.listen_sock:
1441 try:
1442 pkt = self.listen_sock.recv()
1443 except EOFError:
1444 # Socket was closed abruptly. This will likely only
1445 # ever happen when a client socket is passed to the
1446 # automaton (not the case when the automaton is
1447 # listening on a promiscuous conf.L2sniff)
1448 self.listen_sock.close()
1449 # False so that it is still reset by update_sock
1450 self.listen_sock = False # type: ignore
1451 fds.remove(fd)
1452 if self.state.state in self.eofs:
1453 # There is an eof state
1454 eof = self.eofs[self.state.state]
1455 self.debug(2, "Condition EOF [%s] taken" % eof.__name__) # noqa: E501
1456 raise self.eofs[self.state.state](self)
1457 else:
1458 # There isn't. Therefore, it's a closing condition.
1459 raise EOFError("Socket ended arbruptly.")
1460 if self.atmt_session is not None:
1461 # Apply session if provided
1462 pkt = self.atmt_session.process(pkt)
1463 if pkt is not None:
1464 if self.master_filter(pkt):
1465 self.debug(3, "RECVD: %s" % pkt.summary()) # noqa: E501
1466 for rcvcond in self.recv_conditions[self.state.state]: # noqa: E501
1467 self._run_condition(rcvcond, pkt, *state_output) # noqa: E501
1468 else:
1469 self.debug(4, "FILTR: %s" % pkt.summary()) # noqa: E501
1470 else:
1471 self.debug(3, "IOEVENT on %s" % fd.ioname)
1472 for ioevt in self.ioevents[self.state.state]:
1473 if ioevt.atmt_ioname == fd.ioname:
1474 self._run_condition(ioevt, fd, *state_output) # noqa: E501
1476 except ATMT.NewStateRequested as state_req:
1477 self.debug(2, "switching from [%s] to [%s]" % (self.state.state, state_req.state)) # noqa: E501
1478 self.state = state_req
1479 yield state_req
1481 def __repr__(self):
1482 # type: () -> str
1483 return "<Automaton %s [%s]>" % (
1484 self.__class__.__name__,
1485 ["HALTED", "RUNNING"][self.isrunning()]
1486 )
1488 # Public API
1489 def add_interception_points(self, *ipts):
1490 # type: (Any) -> None
1491 for ipt in ipts:
1492 if hasattr(ipt, "atmt_state"):
1493 ipt = ipt.atmt_state
1494 self.interception_points.add(ipt)
1496 def remove_interception_points(self, *ipts):
1497 # type: (Any) -> None
1498 for ipt in ipts:
1499 if hasattr(ipt, "atmt_state"):
1500 ipt = ipt.atmt_state
1501 self.interception_points.discard(ipt)
1503 def add_breakpoints(self, *bps):
1504 # type: (Any) -> None
1505 for bp in bps:
1506 if hasattr(bp, "atmt_state"):
1507 bp = bp.atmt_state
1508 self.breakpoints.add(bp)
1510 def remove_breakpoints(self, *bps):
1511 # type: (Any) -> None
1512 for bp in bps:
1513 if hasattr(bp, "atmt_state"):
1514 bp = bp.atmt_state
1515 self.breakpoints.discard(bp)
1517 def start(self, *args, **kargs):
1518 # type: (Any, Any) -> None
1519 if self.isrunning():
1520 raise ValueError("Already started")
1521 # Start the control thread
1522 self._do_start(*args, **kargs)
1524 def run(self,
1525 resume=None, # type: Optional[Message]
1526 wait=True # type: Optional[bool]
1527 ):
1528 # type: (...) -> Any
1529 if resume is None:
1530 resume = Message(type=_ATMT_Command.RUN)
1531 self.cmdin.send(resume)
1532 if wait:
1533 try:
1534 c = self.cmdout.recv()
1535 if c is None:
1536 return None
1537 except KeyboardInterrupt:
1538 self.cmdin.send(Message(type=_ATMT_Command.FREEZE))
1539 return None
1540 if c.type == _ATMT_Command.END:
1541 return c.result
1542 elif c.type == _ATMT_Command.INTERCEPT:
1543 raise self.InterceptionPoint("packet intercepted", state=c.state.state, packet=c.pkt) # noqa: E501
1544 elif c.type == _ATMT_Command.SINGLESTEP:
1545 raise self.Singlestep("singlestep state=[%s]" % c.state.state, state=c.state.state) # noqa: E501
1546 elif c.type == _ATMT_Command.BREAKPOINT:
1547 raise self.Breakpoint("breakpoint triggered on state [%s]" % c.state.state, state=c.state.state) # noqa: E501
1548 elif c.type == _ATMT_Command.EXCEPTION:
1549 # this code comes from the `six` module (`.reraise()`)
1550 # to raise an exception with specified exc_info.
1551 value = c.exc_info[0]() if c.exc_info[1] is None else c.exc_info[1] # type: ignore # noqa: E501
1552 if value.__traceback__ is not c.exc_info[2]:
1553 raise value.with_traceback(c.exc_info[2])
1554 raise value
1555 return None
1557 def runbg(self, resume=None, wait=False):
1558 # type: (Optional[Message], Optional[bool]) -> None
1559 self.run(resume, wait)
1561 def __next__(self):
1562 # type: () -> Any
1563 return self.run(resume=Message(type=_ATMT_Command.NEXT))
1565 def _flush_inout(self):
1566 # type: () -> None
1567 # Flush command pipes
1568 for cmd in [self.cmdin, self.cmdout]:
1569 cmd.clear()
1571 def destroy(self):
1572 # type: () -> None
1573 """
1574 Destroys a stopped Automaton: this cleanups all opened file descriptors.
1575 Required on PyPy for instance where the garbage collector behaves differently.
1576 """
1577 if not hasattr(self, "started"):
1578 return # was never started.
1579 if self.isrunning():
1580 raise ValueError("Can't close running Automaton ! Call stop() beforehand")
1581 # Close command pipes
1582 self.cmdin.close()
1583 self.cmdout.close()
1584 self._flush_inout()
1585 # Close opened ioins/ioouts
1586 for i in itertools.chain(self.ioin.values(), self.ioout.values()):
1587 if isinstance(i, ObjectPipe):
1588 i.close()
1590 def stop(self, wait=True):
1591 # type: (bool) -> None
1592 try:
1593 self.cmdin.send(Message(type=_ATMT_Command.STOP))
1594 except OSError:
1595 pass
1596 if wait:
1597 with self.started:
1598 self._flush_inout()
1600 def forcestop(self, wait=True):
1601 # type: (bool) -> None
1602 try:
1603 self.cmdin.send(Message(type=_ATMT_Command.FORCESTOP))
1604 except OSError:
1605 pass
1606 if wait:
1607 with self.started:
1608 self._flush_inout()
1610 def restart(self, *args, **kargs):
1611 # type: (Any, Any) -> None
1612 self.stop()
1613 self.start(*args, **kargs)
1615 def accept_packet(self,
1616 pkt=None, # type: Optional[Packet]
1617 wait=False # type: Optional[bool]
1618 ):
1619 # type: (...) -> Any
1620 rsm = Message()
1621 if pkt is None:
1622 rsm.type = _ATMT_Command.ACCEPT
1623 else:
1624 rsm.type = _ATMT_Command.REPLACE
1625 rsm.pkt = pkt
1626 return self.run(resume=rsm, wait=wait)
1628 def reject_packet(self,
1629 wait=False # type: Optional[bool]
1630 ):
1631 # type: (...) -> Any
1632 rsm = Message(type=_ATMT_Command.REJECT)
1633 return self.run(resume=rsm, wait=wait)