1## This file is part of Scapy
2## See http://www.secdev.org/projects/scapy for more informations
3## Copyright (C) Philippe Biondi <phil@secdev.org>
4## Copyright (C) Gabriel Potter <gabriel@potter.fr>
5## This program is published under a GPLv2 license
6
7"""
8Python 2 and 3 link classes.
9"""
10
11from __future__ import absolute_import
12import base64
13import binascii
14
15import scapy.modules.six as six
16
17###########
18# Python3 #
19###########
20
21def cmp_to_key(mycmp):
22    # TODO remove me once all 'key=cmp_to_key(..)' has been fixed in utils6.py, automaton.py
23    """Convert a cmp= function into a key= function.
24    To use with sort()
25
26    e.g: def stg_cmp(a, b):
27            return a == b
28    list.sort(key=cmp_to_key(stg_cmp))
29    """
30    class K(object):
31        def __init__(self, obj, *args):
32            self.obj = obj
33        def __lt__(self, other):
34            return mycmp(self.obj, other.obj) < 0
35        def __gt__(self, other):
36            return mycmp(self.obj, other.obj) > 0
37        def __eq__(self, other):
38            return mycmp(self.obj, other.obj) == 0
39        def __le__(self, other):
40            return mycmp(self.obj, other.obj) <= 0
41        def __ge__(self, other):
42            return mycmp(self.obj, other.obj) >= 0
43        def __ne__(self, other):
44            return mycmp(self.obj, other.obj) != 0
45    return K
46
47def cmp(a, b):
48    """Old Python 2 function"""
49    return (a > b) - (a < b)
50
51
52if six.PY2:
53    def orb(x):
54        """Return ord(x) when necessary."""
55        if isinstance(x, basestring):
56            return ord(x)
57        return x
58else:
59    def orb(x):
60        """Return ord(x) when necessary."""
61        if isinstance(x, (bytes, str)):
62            return ord(x)
63        return x
64
65
66if six.PY2:
67    def raw(x):
68        """Convert a str, a packet to bytes"""
69        if x is None:
70            return None
71        if hasattr(x, "__bytes__"):
72            return x.__bytes__()
73        try:
74            return chr(x)
75        except (ValueError, TypeError):
76            return str(x)
77
78    def plain_str(x):
79        """Convert basic byte objects to str"""
80        return x if isinstance(x, basestring) else str(x)
81
82    def chb(x):
83        """Same than chr() but encode as bytes.
84
85        """
86        if isinstance(x, bytes):
87            return x
88        else:
89            if hasattr(x, "__int__") and not isinstance(x, int):
90                return bytes(chr(int(x)))
91            return bytes(chr(x))
92else:
93    def raw(x):
94        """Convert a str, an int, a list of ints, a packet to bytes"""
95        try:
96            return bytes(x)
97        except TypeError:
98            return bytes(x, encoding="utf8")
99
100    def plain_str(x):
101        """Convert basic byte objects to str"""
102        if isinstance(x, bytes):
103            return x.decode('utf8')
104        return x if isinstance(x, str) else str(x)
105
106    def chb(x):
107        """Same than chr() but encode as bytes.
108
109        """
110        if isinstance(x, bytes):
111            return x
112        else:
113            if hasattr(x, "__int__") and not isinstance(x, int):
114                return bytes([int(x)])
115            return bytes([x])
116
117def bytes_hex(x):
118    """Hexify a str or a bytes object"""
119    return binascii.b2a_hex(raw(x))
120
121def hex_bytes(x):
122    """De-hexify a str or a byte object"""
123    return binascii.a2b_hex(raw(x))
124
125def base64_bytes(x):
126    """Turn base64 into bytes"""
127    if six.PY2:
128        return base64.decodestring(x)
129    return base64.decodebytes(raw(x))
130
131def bytes_base64(x):
132    """Turn bytes into base64"""
133    if six.PY2:
134        return base64.encodestring(x).replace('\n', '')
135    return base64.encodebytes(raw(x)).replace(b'\n', b'')
136