/
dadict.py
166 lines (130 loc) · 4.1 KB
/
dadict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
# SPDX-License-Identifier: GPL-2.0-only
# This file is part of Scapy
# See https://scapy.net/ for more information
# Copyright (C) Philippe Biondi <phil@secdev.org>
"""
Direct Access dictionary.
"""
from scapy.error import Scapy_Exception
from scapy.compat import plain_str
# Typing
from typing import (
Any,
Dict,
Generic,
Iterator,
List,
Tuple,
Type,
TypeVar,
Union,
)
from scapy.compat import Self
###############################
# Direct Access dictionary #
###############################
def fixname(x):
# type: (Union[bytes, str]) -> str
"""
Modifies a string to make sure it can be used as an attribute name.
"""
x = plain_str(x)
if x and str(x[0]) in "0123456789":
x = "n_" + x
return x.translate(
"________________________________________________"
"0123456789_______ABCDEFGHIJKLMNOPQRSTUVWXYZ______"
"abcdefghijklmnopqrstuvwxyz____________________________"
"______________________________________________________"
"___________________________________________________"
)
class DADict_Exception(Scapy_Exception):
pass
_K = TypeVar('_K') # Key type
_V = TypeVar('_V') # Value type
class DADict(Generic[_K, _V]):
"""
Direct Access Dictionary
This acts like a dict, but it provides a direct attribute access
to its keys through its values. This is used to store protocols,
manuf...
For instance, scapy fields will use a DADict as an enum::
ETHER_TYPES[2048] -> IPv4
Whereas humans can access::
ETHER_TYPES.IPv4 -> 2048
"""
__slots__ = ["_name", "d"]
def __init__(self, _name="DADict", **kargs):
# type: (str, **Any) -> None
self._name = _name
self.d = {} # type: Dict[_K, _V]
self.update(kargs) # type: ignore
def ident(self, v):
# type: (_V) -> str
"""
Return value that is used as key for the direct access
"""
if isinstance(v, (str, bytes)):
return fixname(v)
return "unknown"
def update(self, *args, **kwargs):
# type: (*Dict[_K, _V], **Dict[_K, _V]) -> None
for k, v in dict(*args, **kwargs).items():
self[k] = v # type: ignore
def iterkeys(self):
# type: () -> Iterator[_K]
for x in self.d:
if not isinstance(x, str) or x[0] != "_":
yield x
def keys(self):
# type: () -> List[_K]
return list(self.iterkeys())
def __iter__(self):
# type: () -> Iterator[_K]
return self.iterkeys()
def itervalues(self):
# type: () -> Iterator[_V]
return self.d.values() # type: ignore
def values(self):
# type: () -> List[_V]
return list(self.itervalues())
def _show(self):
# type: () -> None
for k in self.iterkeys():
print("%10s = %r" % (k, self[k]))
def __repr__(self):
# type: () -> str
return "<%s - %s elements>" % (self._name, len(self))
def __getitem__(self, attr):
# type: (_K) -> _V
return self.d[attr]
def __setitem__(self, attr, val):
# type: (_K, _V) -> None
self.d[attr] = val
def __len__(self):
# type: () -> int
return len(self.d)
def __nonzero__(self):
# type: () -> bool
# Always has at least its name
return len(self) > 1
__bool__ = __nonzero__
def __getattr__(self, attr):
# type: (str) -> _K
try:
return object.__getattribute__(self, attr) # type: ignore
except AttributeError:
for k, v in self.d.items():
if self.ident(v) == attr:
return k
raise AttributeError
def __dir__(self):
# type: () -> List[str]
return [self.ident(x) for x in self.itervalues()]
def __reduce__(self):
# type: () -> Tuple[Type[Self], Tuple[str], Tuple[Dict[_K, _V]]]
return (self.__class__, (self._name,), (self.d,))
def __setstate__(self, state):
# type: (Tuple[Dict[_K, _V]]) -> Self
self.d.update(state[0])
return self