1import re
2import warnings
3from dataclasses import dataclass, is_dataclass, fields
4from typing import ClassVar, Any, Dict, Type, Set, List, TYPE_CHECKING, Sequence
5
6import mitmproxy.flow
7
8if TYPE_CHECKING:
9    import mitmproxy.addonmanager
10    import mitmproxy.log
11
12
13class Hook:
14    name: ClassVar[str]
15
16    def args(self) -> List[Any]:
17        args = []
18        for field in fields(self):
19            args.append(getattr(self, field.name))
20        return args
21
22    def __new__(cls, *args, **kwargs):
23        if cls is Hook:
24            raise TypeError("Hook may not be instantiated directly.")
25        if not is_dataclass(cls):
26            raise TypeError("Subclass is not a dataclass.")
27        return super().__new__(cls)
28
29    def __init_subclass__(cls, **kwargs):
30        # initialize .name attribute. HttpRequestHook -> http_request
31        if cls.__dict__.get("name", None) is None:
32            name = cls.__name__.replace("Hook", "")
33            cls.name = re.sub('(?!^)([A-Z]+)', r'_\1', name).lower()
34        if cls.name in all_hooks:
35            other = all_hooks[cls.name]
36            warnings.warn(f"Two conflicting event classes for {cls.name}: {cls} and {other}", RuntimeWarning)
37        if cls.name == "":
38            return  # don't register Hook class.
39        all_hooks[cls.name] = cls
40
41        # define a custom hash and __eq__ function so that events are hashable and not comparable.
42        cls.__hash__ = object.__hash__
43        cls.__eq__ = object.__eq__
44
45
46all_hooks: Dict[str, Type[Hook]] = {}
47
48
49@dataclass
50class ConfigureHook(Hook):
51    """
52    Called when configuration changes. The updated argument is a
53    set-like object containing the keys of all changed options. This
54    event is called during startup with all options in the updated set.
55    """
56    updated: Set[str]
57
58
59@dataclass
60class DoneHook(Hook):
61    """
62    Called when the addon shuts down, either by being removed from
63    the mitmproxy instance, or when mitmproxy itself shuts down. On
64    shutdown, this event is called after the event loop is
65    terminated, guaranteeing that it will be the final event an addon
66    sees. Note that log handlers are shut down at this point, so
67    calls to log functions will produce no output.
68    """
69
70
71@dataclass
72class RunningHook(Hook):
73    """
74    Called when the proxy is completely up and running. At this point,
75    you can expect the proxy to be bound to a port, and all addons to be
76    loaded.
77    """
78
79
80@dataclass
81class UpdateHook(Hook):
82    """
83    Update is called when one or more flow objects have been modified,
84    usually from a different addon.
85    """
86    flows: Sequence[mitmproxy.flow.Flow]
87