1# Copyright (c) 2013-2022 Intel Corporation.
2#
3# Redistribution and use in source and binary forms, with or without
4# modification, are permitted provided that the following conditions are met:
5#
6#     * Redistributions of source code must retain the above copyright notice,
7#       this list of conditions and the following disclaimer.
8#     * Redistributions in binary form must reproduce the above copyright notice,
9#       this list of conditions and the following disclaimer in the documentation
10#       and/or other materials provided with the distribution.
11#     * Neither the name of Intel Corporation nor the names of its contributors
12#       may be used to endorse or promote products derived from this software
13#       without specific prior written permission.
14#
15# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
16# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
17# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
18# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
19# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
20# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
21# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
22# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
23# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
24# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
25
26"""unpack module."""
27
28from collections import OrderedDict
29import struct
30
31class UnpackError(Exception):
32    pass
33
34class Unpackable(object):
35    def __init__(self, data, offset=0, size=None):
36        self.data = data
37        data_size = len(data)
38        if offset > data_size:
39            raise UnpackError("Unpackable.__init__: offset={} but len(data)={}".format(offset, data_size))
40        self.offset = offset
41        if size is None:
42            self.size = data_size
43        else:
44            self.size = offset + size
45            if self.size > data_size:
46                raise UnpackError("Unpackable.__init__: offset+size={} but len(data)={}".format(self.size, data_size))
47
48    def _check_unpack(self, size):
49        if self.offset + size > self.size:
50            raise UnpackError("Unpackable: Attempted to unpack {} bytes, but only {} bytes remaining".format(size, self.size - self.offset))
51
52    def skip(self, size):
53        self._check_unpack(size)
54        self.offset += size
55
56    def unpack(self, fmt):
57        try:
58            l = struct.calcsize(fmt)
59            self._check_unpack(l)
60            value = struct.unpack_from(fmt, self.data, self.offset)
61            self.offset += l
62            return value
63        except struct.error as e:
64            raise UnpackError("Unpackable.unpack: " + str(e))
65
66    def unpack_one(self, fmt):
67        return self.unpack(fmt)[0]
68
69    def unpack_peek(self, fmt):
70        try:
71            l = struct.calcsize(fmt)
72            self._check_unpack(l)
73            return struct.unpack_from(fmt, self.data, self.offset)
74        except struct.error as e:
75            raise UnpackError("Unpackable.unpack_peek: " + str(e))
76
77    def unpack_peek_one(self, fmt):
78        return self.unpack_peek(fmt)[0]
79
80    def unpack_peek_raw(self, size):
81        """Peek at the specified number of bytes as a str"""
82        self._check_unpack(size)
83        return self.data[self.offset:self.offset+size]
84
85    def unpack_peek_rest(self):
86        """Peek at the remainder of the unpackable as a str"""
87        return self.data[self.offset:self.size]
88
89    def unpack_raw(self, size):
90        """Unpack the specified number of bytes as a str"""
91        self._check_unpack(size)
92        old_offset = self.offset
93        self.offset += size
94        return self.data[old_offset:self.offset]
95
96    def unpack_rest(self):
97        """Return the remainder of the unpackable as a str"""
98        offset = self.offset
99        self.offset = self.size
100        return self.data[offset:self.size]
101
102    def unpack_unpackable(self, size):
103        """Unpack the specified number of bytes as an Unpackable"""
104        u = Unpackable(self.data, self.offset, size)
105        self.offset += size
106        return u
107
108    def at_end(self):
109        return self.offset == self.size
110
111class StructError(Exception):
112    pass
113
114class Struct(object):
115    def __init__(self):
116        self.fields = OrderedDict()
117
118    @classmethod
119    def unpack(cls, u):
120        s = cls()
121        for field in cls._unpack(u):
122            s.add_field(*field)
123        return s
124
125    def add_field(self, name, value, fmt=None):
126        if hasattr(self, name):
127            raise StructError("Internal error: Duplicate Struct field name {}".format(name))
128        if fmt is None:
129            if isinstance(value, int) and not isinstance(value, bool):
130                fmt = "{:#x}".format
131            else:
132                fmt = "{!r}".format
133        elif isinstance(fmt, str):
134            fmt = fmt.format
135        elif not callable(fmt):
136            raise StructError("Internal error: Expected a format string or callable, but got: {}".format(fmt))
137        setattr(self, name, value)
138        self.fields[name] = fmt
139
140    def format_field(self, name):
141        return self.fields[name](getattr(self, name))
142
143    def __repr__(self):
144        return "{}({})".format(self.__class__.__name__, ", ".join("{}={}".format(k, self.format_field(k)) for k in self.fields.keys()))
145
146    def __iter__(self):
147        return (getattr(self, k) for k in self.fields.keys())
148
149    def __eq__(self, other):
150        if type(self) is not type(other):
151            return NotImplemented
152        return self.fields.keys() == other.fields.keys() and all(getattr(self, name) == getattr(other, name) for name in self.fields.keys())
153
154    def __ne__(self, other):
155        return not self == other
156
157    def __hash__(self):
158        return hash(tuple((name, getattr(self, name)) for name in self.fields.keys()))
159
160def format_each(fmt_one):
161    def f(it):
162        return "({})".format(", ".join(fmt_one.format(i) for i in it))
163    return f
164
165format_each_hex = format_each("{:#x}")
166
167def format_table(fmt, table, default='Reserved'):
168    def f(value):
169        return "{} ({})".format(fmt.format(value), table.get(value, default))
170    return f
171
172def format_function(fmt, function):
173    def f(value):
174        return "{} ({})".format(fmt.format(value), function(value))
175    return f
176
177def reserved_None(fmt="{!r}"):
178    def f(value):
179        if value is None:
180            return "Reserved"
181        return fmt.format(value)
182    return f
183
184def unpack_all(u, structs, *args):
185    """Keep constructing structs from the unpackable u until it runs out of data.
186
187    structs should consist of a list of Struct subclasses to be tried in order.
188    Each of them should return None from their constructor if they're not the
189    correct type to unpack the next chunk of data.  Any catch-all generic
190    structure should apepar last in the list.  Raises a StructError if no
191    struct matches."""
192    def _substructs():
193        while not u.at_end():
194            for s in structs:
195                temp = s(u, *args)
196                if temp is not None:
197                    yield temp
198                    break
199            else:
200                raise StructError("Internal error: unable to unpack any structure at byte {} of unpackable".format(u.offset))
201    return tuple(_substructs())
202