1# This file is part of Scapy
2# See http://www.secdev.org/projects/scapy for more information
3# Copyright (C) Gabriel Potter <gabriel@potter.fr>
4# This program is published under a GPLv2 license
5
6"""
7Python 2 and 3 link classes.
8"""
9
10from __future__ import absolute_import
11import base64
12import binascii
13import collections
14import gzip
15import socket
16import struct
17import sys
18
19import scapy.modules.six as six
20
21# Very important: will issue typing errors otherwise
22__all__ = [
23    # typing
24    'Any',
25    'AnyStr',
26    'Callable',
27    'DefaultDict',
28    'Dict',
29    'Generic',
30    'IO',
31    'Iterable',
32    'Iterator',
33    'List',
34    'Literal',
35    'NamedTuple',
36    'NewType',
37    'NoReturn',
38    'Optional',
39    'Pattern',
40    'Sequence',
41    'Set',
42    'Sized',
43    'Tuple',
44    'Type',
45    'TypeVar',
46    'Union',
47    'cast',
48    'overload',
49    'FAKE_TYPING',
50    'TYPE_CHECKING',
51    # compat
52    'AddressFamily',
53    'base64_bytes',
54    'bytes_base64',
55    'bytes_encode',
56    'bytes_hex',
57    'chb',
58    'gzip_compress',
59    'gzip_decompress',
60    'hex_bytes',
61    'lambda_tuple_converter',
62    'orb',
63    'plain_str',
64    'raw',
65]
66
67# Typing compatibility
68
69# Note:
70# supporting typing on multiple python versions is a nightmare.
71# Since Python 3.7, Generic is a type instead of a metaclass,
72# therefore we can't support both at the same time. Our strategy
73# is to only use the typing module if the Python version is >= 3.7
74# and use totally fake replacements otherwise.
75# HOWEVER, when using the fake ones, to emulate stub Generic
76# fields (e.g. _PacketField[str]) we need to add a fake
77# __getitem__ to Field_metaclass
78
79try:
80    import typing  # noqa: F401
81    from typing import TYPE_CHECKING
82    if sys.version_info[0:2] <= (3, 6):
83        # Generic is messed up before Python 3.7
84        # https://github.com/python/typing/issues/449
85        raise ImportError
86    FAKE_TYPING = False
87except ImportError:
88    FAKE_TYPING = True
89    TYPE_CHECKING = False
90
91# Import or create fake types
92
93
94def _FakeType(name, cls=object):
95    # type: (str, Optional[type]) -> Any
96    class _FT(object):
97        def __init__(self, name):
98            # type: (str) -> None
99            self.name = name
100
101        # make the objects subscriptable indefinetly
102        def __getitem__(self, item):  # type: ignore
103            return cls
104
105        def __call__(self, *args, **kargs):
106            # type: (*Any, **Any) -> Any
107            if isinstance(args[0], str):
108                self.name = args[0]
109            return self
110
111        def __repr__(self):
112            # type: () -> str
113            return "<Fake typing.%s>" % self.name
114    return _FT(name)
115
116
117if not FAKE_TYPING:
118    # Only required if using mypy-lang for static typing
119    from typing import (
120        Any,
121        AnyStr,
122        Callable,
123        DefaultDict,
124        Dict,
125        Generic,
126        Iterable,
127        Iterator,
128        IO,
129        List,
130        NewType,
131        NoReturn,
132        Optional,
133        Pattern,
134        Sequence,
135        Set,
136        Sized,
137        Tuple,
138        Type,
139        TypeVar,
140        Union,
141        cast,
142        overload,
143    )
144else:
145    # Let's be creative and make some fake ones.
146    def cast(_type, obj):  # type: ignore
147        return obj
148
149    Any = _FakeType("Any")
150    AnyStr = _FakeType("AnyStr")  # type: ignore
151    Callable = _FakeType("Callable")
152    DefaultDict = _FakeType("DefaultDict",  # type: ignore
153                            collections.defaultdict)
154    Dict = _FakeType("Dict", dict)  # type: ignore
155    Generic = _FakeType("Generic")
156    Iterable = _FakeType("Iterable")  # type: ignore
157    Iterator = _FakeType("Iterator")  # type: ignore
158    IO = _FakeType("IO")  # type: ignore
159    List = _FakeType("List", list)  # type: ignore
160    NewType = _FakeType("NewType")
161    NoReturn = _FakeType("NoReturn")  # type: ignore
162    Optional = _FakeType("Optional")
163    Pattern = _FakeType("Pattern")  # type: ignore
164    Sequence = _FakeType("Sequence")  # type: ignore
165    Set = _FakeType("Set", set)  # type: ignore
166    Sequence = _FakeType("Sequence", list)  # type: ignore
167    Tuple = _FakeType("Tuple")
168    Type = _FakeType("Type", type)
169    TypeVar = _FakeType("TypeVar")  # type: ignore
170    Union = _FakeType("Union")
171
172    class Sized(object):  # type: ignore
173        pass
174
175    overload = lambda x: x
176
177
178# Broken < Python 3.7
179if sys.version_info >= (3, 7):
180    from typing import NamedTuple
181else:
182    # Hack for Python < 3.7 - Implement NamedTuple pickling
183    def _unpickleNamedTuple(name, len_params, *args):
184        return collections.namedtuple(
185            name,
186            args[:len_params]
187        )(*args[len_params:])
188
189    def NamedTuple(name, params):
190        tup_params = tuple(x[0] for x in params)
191        cls = collections.namedtuple(name, tup_params)
192
193        class _NT(cls):
194            def __reduce__(self):
195                """Used by pickling methods"""
196                return (_unpickleNamedTuple,
197                        (name, len(tup_params)) + tup_params + tuple(self))
198        _NT.__name__ = cls.__name__
199        return _NT
200
201# Python 3.8 Only
202if sys.version_info >= (3, 8):
203    from typing import Literal
204else:
205    Literal = _FakeType("Literal")
206
207# Python 3.4
208if sys.version_info >= (3, 4):
209    from socket import AddressFamily
210else:
211    class AddressFamily:
212        AF_INET = socket.AF_INET
213        AF_INET6 = socket.AF_INET6
214
215
216class _Generic_metaclass(type):
217    if FAKE_TYPING:
218        def __getitem__(self, typ):
219            # type: (Any) -> Any
220            return self
221
222
223###########
224# Python3 #
225###########
226
227# https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators
228DecoratorCallable = TypeVar("DecoratorCallable", bound=Callable[..., Any])
229
230
231def lambda_tuple_converter(func):
232    # type: (DecoratorCallable) -> DecoratorCallable
233    """
234    Converts a Python 2 function as
235      lambda (x,y): x + y
236    In the Python 3 format:
237      lambda x,y : x + y
238    """
239    if func is not None and func.__code__.co_argcount == 1:
240        return lambda *args: func(  # type: ignore
241            args[0] if len(args) == 1 else args
242        )
243    else:
244        return func
245
246
247# This is ugly, but we don't want to move raw() out of compat.py
248# and it makes it much clearer
249if TYPE_CHECKING:
250    from scapy.packet import Packet
251
252
253if six.PY2:
254    bytes_encode = plain_str = str  # type: Callable[[Any], bytes]
255    orb = ord  # type: Callable[[bytes], int]
256
257    def chb(x):
258        # type: (int) -> bytes
259        if isinstance(x, str):
260            return x
261        return chr(x)
262
263    def raw(x):
264        # type: (Union[Packet]) -> bytes
265        """
266        Builds a packet and returns its bytes representation.
267        This function is and will always be cross-version compatible
268        """
269        if hasattr(x, "__bytes__"):
270            return x.__bytes__()
271        return bytes(x)
272else:
273    def raw(x):
274        # type: (Union[Packet]) -> bytes
275        """
276        Builds a packet and returns its bytes representation.
277        This function is and will always be cross-version compatible
278        """
279        return bytes(x)
280
281    def bytes_encode(x):
282        # type: (Any) -> bytes
283        """Ensure that the given object is bytes.
284        If the parameter is a packet, raw() should be preferred.
285        """
286        if isinstance(x, str):
287            return x.encode()
288        return bytes(x)
289
290    if sys.version_info[0:2] <= (3, 4):
291        def plain_str(x):
292            # type: (AnyStr) -> str
293            """Convert basic byte objects to str"""
294            if isinstance(x, bytes):
295                return x.decode(errors="ignore")
296            return str(x)
297    else:
298        # Python 3.5+
299        def plain_str(x):
300            # type: (Any) -> str
301            """Convert basic byte objects to str"""
302            if isinstance(x, bytes):
303                return x.decode(errors="backslashreplace")
304            return str(x)
305
306    def chb(x):
307        # type: (int) -> bytes
308        """Same than chr() but encode as bytes."""
309        return struct.pack("!B", x)
310
311    def orb(x):
312        # type: (Union[int, str, bytes]) -> int
313        """Return ord(x) when not already an int."""
314        if isinstance(x, int):
315            return x
316        return ord(x)
317
318
319def bytes_hex(x):
320    # type: (AnyStr) -> bytes
321    """Hexify a str or a bytes object"""
322    return binascii.b2a_hex(bytes_encode(x))
323
324
325def hex_bytes(x):
326    # type: (AnyStr) -> bytes
327    """De-hexify a str or a byte object"""
328    return binascii.a2b_hex(bytes_encode(x))
329
330
331def base64_bytes(x):
332    # type: (AnyStr) -> bytes
333    """Turn base64 into bytes"""
334    if six.PY2:
335        return base64.decodestring(x)  # type: ignore
336    return base64.decodebytes(bytes_encode(x))
337
338
339def bytes_base64(x):
340    # type: (AnyStr) -> bytes
341    """Turn bytes into base64"""
342    if six.PY2:
343        return base64.encodestring(x).replace('\n', '')  # type: ignore
344    return base64.encodebytes(bytes_encode(x)).replace(b'\n', b'')
345
346
347if six.PY2:
348    import cgi
349    html_escape = cgi.escape
350else:
351    import html
352    html_escape = html.escape
353
354
355if six.PY2:
356    from StringIO import StringIO
357
358    def gzip_decompress(x):
359        # type: (AnyStr) -> bytes
360        """Decompress using gzip"""
361        with gzip.GzipFile(fileobj=StringIO(x), mode='rb') as fdesc:
362            return fdesc.read()
363
364    def gzip_compress(x):
365        # type: (AnyStr) -> bytes
366        """Compress using gzip"""
367        buf = StringIO()
368        with gzip.GzipFile(fileobj=buf, mode='wb') as fdesc:
369            fdesc.write(x)
370        return buf.getvalue()
371else:
372    gzip_decompress = gzip.decompress
373    gzip_compress = gzip.compress
374