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