1#!/usr/bin/env python3
2#
3#  Copyright (c) 2019, The OpenThread Authors.
4#  All rights reserved.
5#
6#  Redistribution and use in source and binary forms, with or without
7#  modification, are permitted provided that the following conditions are met:
8#  1. Redistributions of source code must retain the above copyright
9#     notice, this list of conditions and the following disclaimer.
10#  2. Redistributions in binary form must reproduce the above copyright
11#     notice, this list of conditions and the following disclaimer in the
12#     documentation and/or other materials provided with the distribution.
13#  3. Neither the name of the copyright holder nor the
14#     names of its contributors may be used to endorse or promote products
15#     derived from this software without specific prior written permission.
16#
17#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
18#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
21#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
22#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
23#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
24#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
25#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
26#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
27#  POSSIBILITY OF SUCH DAMAGE.
28#
29import sys
30from typing import Union, Any
31
32
33class Bytes(bytearray):
34    """Bytes represents a byte array which is able to handle strings of flexible formats"""
35
36    def __init__(self, s: Union[str, bytearray, 'Bytes', Any]):
37        if isinstance(s, str):
38            try:
39                s = Bytes._parse_compact(s)
40            except ValueError:
41                try:
42                    s = Bytes._parse_octets(s)
43                except ValueError:
44                    s = Bytes._parse_hextets(s)
45
46        super().__init__(s)
47
48    def __hash__(self):
49        return hash(bytes(self))
50
51    def __repr__(self):
52        return '%s(%r)' % (self.__class__.__name__, self.format_compact())
53
54    def format_compact(self) -> str:
55        """
56        Converts the Bytes to a compact string (without ":").
57        """
58        return ''.join('%02x' % b for b in self)
59
60    def format_octets(self) -> str:
61        """
62        Converts the Bytes to a string of octets separated by ":".
63        """
64        return ':'.join('%02x' % b for b in self)
65
66    def format_hextets(self) -> str:
67        """
68        Converts the Bytes to a string of hextets separated by ":"
69        """
70        assert len(self) % 2 == 0, self.format_octets()
71        return ':'.join('%04x' % (self[i] * 256 + self[i + 1]) for i in range(0, len(self), 2))
72
73    __str__ = format_octets
74
75    @staticmethod
76    def _parse_compact(s: str) -> bytearray:
77        try:
78            assert len(s) % 2 == 0
79            return bytearray(int(s[i:i + 2], 16) for i in range(0, len(s), 2))
80        except Exception:
81            raise ValueError(s)
82
83    @staticmethod
84    def _parse_octets(s: str) -> bytearray:
85        try:
86            assert len(s) % 3 == 2 or not s
87            if not s:
88                return bytearray(b"")
89
90            return bytearray(int(x, 16) for x in s.split(':'))
91        except Exception:
92            raise ValueError(s)
93
94    @staticmethod
95    def _parse_hextets(s) -> bytearray:
96        try:
97            assert len(s) % 5 == 4 or not s
98            if not s:
99                return bytearray(b"")
100
101            return bytearray(int(x[i:i + 2], 16) for x in s.split(':') for i in (0, 2))
102        except Exception:
103            raise ValueError(s)
104
105    def __getitem__(self, item) -> Union['Bytes', int]:
106        """
107        Get self[item].
108
109        :param item: index or slice to retrieve
110        :return: the byte value at specified index or sub `Bytes` if item is slice
111        """
112        x = super().__getitem__(item)
113        if isinstance(x, bytearray):
114            return Bytes(x)
115        else:
116            return x
117
118    def __eq__(self, other: Union[str, 'Bytes']):
119        """
120        Check if bytes is equal to other.
121        """
122        if other is None:
123            return False
124        elif not isinstance(other, Bytes):
125            other = self.__class__(other)
126
127        eq = super().__eq__(other)
128        print("[%r %s %r]" % (self, "==" if eq else "!=", other), file=sys.stderr)
129        return eq
130
131
132if __name__ == '__main__':
133    # some simple tests
134    x = Bytes(b"\x01\x02\x03\x04")
135    assert eval(repr(x)) == x, repr(x)  # representation of Bytes should be able to be evaluated back
136    assert x == str(x), (x, str(x))
137
138    assert x.format_compact() == "01020304", x.format_compact()
139    assert x.format_octets() == "01:02:03:04", x.format_octets()
140    assert x.format_hextets() == "0102:0304", x.format_hextets()
141
142    assert Bytes._parse_compact("") == Bytes(b"")
143    assert Bytes._parse_compact('01020304') == x
144
145    assert Bytes._parse_octets("") == Bytes(b"")
146    assert Bytes._parse_octets('01:02:03:04') == x
147
148    assert Bytes._parse_hextets("") == Bytes(b"")
149    assert Bytes._parse_hextets('0102:0304') == x
150
151    assert isinstance(x[:2], Bytes)
152    assert isinstance(x[-2:], Bytes)
153    assert x[:2] == Bytes(b'\x01\x02')
154    assert x[-2:] == Bytes(b'\x03\x04')
155
156    # should also parse string formats
157    assert Bytes("01020304") == Bytes(b"\x01\x02\x03\x04")
158    assert Bytes("01:02:03:04") == Bytes(b"\x01\x02\x03\x04")
159    assert Bytes("0102:0304") == Bytes(b"\x01\x02\x03\x04")
160