src/third_party/pyelftools/elftools/construct/core.py

   1 from struct import Struct as Packer
   2
   3 from .lib.py3compat import BytesIO, advance_iterator, bchr
   4 from .lib import Container, ListContainer, LazyContainer
   5
   6
   7 #===============================================================================
   8 # exceptions
   9 #===============================================================================
  10 class ConstructError(Exception):
  11     __slots__ = []
  12 class FieldError(ConstructError):
  13     __slots__ = []
  14 class SizeofError(ConstructError):
  15     __slots__ = []
  16 class AdaptationError(ConstructError):
  17     __slots__ = []
  18 class ArrayError(ConstructError):
  19     __slots__ = []
  20 class RangeError(ConstructError):
  21     __slots__ = []
  22 class SwitchError(ConstructError):
  23     __slots__ = []
  24 class SelectError(ConstructError):
  25     __slots__ = []
  26 class TerminatorError(ConstructError):
  27     __slots__ = []
  28
  29 #===============================================================================
  30 # abstract constructs
  31 #===============================================================================
  32 class Construct(object):
  33     """
  34     The mother of all constructs.
  35
  36     This object is generally not directly instantiated, and it does not
  37     directly implement parsing and building, so it is largely only of interest
  38     to subclass implementors.
  39
  40     The external user API:
  41
  42      * parse()
  43      * parse_stream()
  44      * build()
  45      * build_stream()
  46      * sizeof()
  47
  48     Subclass authors should not override the external methods. Instead,
  49     another API is available:
  50
  51      * _parse()
  52      * _build()
  53      * _sizeof()
  54
  55     There is also a flag API:
  56
  57      * _set_flag()
  58      * _clear_flag()
  59      * _inherit_flags()
  60      * _is_flag()
  61
  62     And stateful copying:
  63
  64      * __getstate__()
  65      * __setstate__()
  66
  67     Attributes and Inheritance
  68     ==========================
  69
  70     All constructs have a name and flags. The name is used for naming struct
  71     members and context dictionaries. Note that the name can either be a
  72     string, or None if the name is not needed. A single underscore ("_") is a
  73     reserved name, and so are names starting with a less-than character ("<").
  74     The name should be descriptive, short, and valid as a Python identifier,
  75     although these rules are not enforced.
  76
  77     The flags specify additional behavioral information about this construct.
  78     Flags are used by enclosing constructs to determine a proper course of
  79     action. Flags are inherited by default, from inner subconstructs to outer
  80     constructs. The enclosing construct may set new flags or clear existing
  81     ones, as necessary.
  82
  83     For example, if FLAG_COPY_CONTEXT is set, repeaters will pass a copy of
  84     the context for each iteration, which is necessary for OnDemand parsing.
  85     """
  86
  87     FLAG_COPY_CONTEXT          = 0x0001
  88     FLAG_DYNAMIC               = 0x0002
  89     FLAG_EMBED                 = 0x0004
  90     FLAG_NESTING               = 0x0008
  91
  92     __slots__ = ["name", "conflags"]
  93     def __init__(self, name, flags = 0):
  94         if name is not None:
  95             if type(name) is not str:
  96                 raise TypeError("name must be a string or None", name)
  97             if name == "_" or name.startswith("<"):
  98                 raise ValueError("reserved name", name)
  99         self.name = name
 100         self.conflags = flags
 101
 102     def __repr__(self):
 103         return "%s(%r)" % (self.__class__.__name__, self.name)
 104
 105     def _set_flag(self, flag):
 106         """
 107         Set the given flag or flags.
 108
 109         :param int flag: flag to set; may be OR'd combination of flags
 110         """
 111
 112         self.conflags |= flag
 113
 114     def _clear_flag(self, flag):
 115         """
 116         Clear the given flag or flags.
 117
 118         :param int flag: flag to clear; may be OR'd combination of flags
 119         """
 120
 121         self.conflags &= ~flag
 122
 123     def _inherit_flags(self, *subcons):
 124         """
 125         Pull flags from subconstructs.
 126         """
 127
 128         for sc in subcons:
 129             self._set_flag(sc.conflags)
 130
 131     def _is_flag(self, flag):
 132         """
 133         Check whether a given flag is set.
 134
 135         :param int flag: flag to check
 136         """
 137
 138         return bool(self.conflags & flag)
 139
 140     def __getstate__(self):
 141         """
 142         Obtain a dictionary representing this construct's state.
 143         """
 144
 145         attrs = {}
 146         if hasattr(self, "__dict__"):
 147             attrs.update(self.__dict__)
 148         slots = []
 149         c = self.__class__
 150         while c is not None:
 151             if hasattr(c, "__slots__"):
 152                 slots.extend(c.__slots__)
 153             c = c.__base__
 154         for name in slots:
 155             if hasattr(self, name):
 156                 attrs[name] = getattr(self, name)
 157         return attrs
 158
 159     def __setstate__(self, attrs):
 160         """
 161         Set this construct's state to a given state.
 162         """
 163         for name, value in attrs.items():
 164             setattr(self, name, value)
 165
 166     def __copy__(self):
 167         """returns a copy of this construct"""
 168         self2 = object.__new__(self.__class__)
 169         self2.__setstate__(self.__getstate__())
 170         return self2
 171
 172     def parse(self, data):
 173         """
 174         Parse an in-memory buffer.
 175
 176         Strings, buffers, memoryviews, and other complete buffers can be
 177         parsed with this method.
 178         """
 179
 180         return self.parse_stream(BytesIO(data))
 181
 182     def parse_stream(self, stream):
 183         """
 184         Parse a stream.
 185
 186         Files, pipes, sockets, and other streaming sources of data are handled
 187         by this method.
 188         """
 189
 190         return self._parse(stream, Container())
 191
 192     def _parse(self, stream, context):
 193         """
 194         Override me in your subclass.
 195         """
 196
 197         raise NotImplementedError()
 198
 199     def build(self, obj):
 200         """
 201         Build an object in memory.
 202         """
 203         stream = BytesIO()
 204         self.build_stream(obj, stream)
 205         return stream.getvalue()
 206
 207     def build_stream(self, obj, stream):
 208         """
 209         Build an object directly into a stream.
 210         """
 211         self._build(obj, stream, Container())
 212
 213     def _build(self, obj, stream, context):
 214         """
 215         Override me in your subclass.
 216         """
 217
 218         raise NotImplementedError()
 219
 220     def sizeof(self, context=None):
 221         """
 222         Calculate the size of this object, optionally using a context.
 223
 224         Some constructs have no fixed size and can only know their size for a
 225         given hunk of data; these constructs will raise an error if they are
 226         not passed a context.
 227
 228         :param ``Container`` context: contextual data
 229
 230         :returns: int of the length of this construct
 231         :raises SizeofError: the size could not be determined
 232         """
 233
 234         if context is None:
 235             context = Container()
 236         try:
 237             return self._sizeof(context)
 238         except Exception as e:
 239             raise SizeofError(e)
 240
 241     def _sizeof(self, context):
 242         """
 243         Override me in your subclass.
 244         """
 245
 246         raise SizeofError("Raw Constructs have no size!")
 247
 248 class Subconstruct(Construct):
 249     """
 250     Abstract subconstruct (wraps an inner construct, inheriting its
 251     name and flags).
 252
 253     Parameters:
 254     * subcon - the construct to wrap
 255     """
 256     __slots__ = ["subcon"]
 257     def __init__(self, subcon):
 258         Construct.__init__(self, subcon.name, subcon.conflags)
 259         self.subcon = subcon
 260     def _parse(self, stream, context):
 261         return self.subcon._parse(stream, context)
 262     def _build(self, obj, stream, context):
 263         self.subcon._build(obj, stream, context)
 264     def _sizeof(self, context):
 265         return self.subcon._sizeof(context)
 266
 267 class Adapter(Subconstruct):
 268     """
 269     Abstract adapter: calls _decode for parsing and _encode for building.
 270
 271     Parameters:
 272     * subcon - the construct to wrap
 273     """
 274     __slots__ = []
 275     def _parse(self, stream, context):
 276         return self._decode(self.subcon._parse(stream, context), context)
 277     def _build(self, obj, stream, context):
 278         self.subcon._build(self._encode(obj, context), stream, context)
 279     def _decode(self, obj, context):
 280         raise NotImplementedError()
 281     def _encode(self, obj, context):
 282         raise NotImplementedError()
 283
 284
 285 #===============================================================================
 286 # Fields
 287 #===============================================================================
 288 def _read_stream(stream, length):
 289     if length < 0:
 290         raise ValueError("length must be >= 0", length)
 291     data = stream.read(length)
 292     if len(data) != length:
 293         raise FieldError("expected %d, found %d" % (length, len(data)))
 294     return data
 295
 296 def _write_stream(stream, length, data):
 297     if length < 0:
 298         raise ValueError("length must be >= 0", length)
 299     if len(data) != length:
 300         raise FieldError("expected %d, found %d" % (length, len(data)))
 301     stream.write(data)
 302
 303 class StaticField(Construct):
 304     """
 305     A fixed-size byte field.
 306
 307     :param str name: field name
 308     :param int length: number of bytes in the field
 309     """
 310
 311     __slots__ = ["length"]
 312     def __init__(self, name, length):
 313         Construct.__init__(self, name)
 314         self.length = length
 315     def _parse(self, stream, context):
 316         return _read_stream(stream, self.length)
 317     def _build(self, obj, stream, context):
 318         _write_stream(stream, self.length, obj)
 319     def _sizeof(self, context):
 320         return self.length
 321
 322 class FormatField(StaticField):
 323     """
 324     A field that uses ``struct`` to pack and unpack data.
 325
 326     See ``struct`` documentation for instructions on crafting format strings.
 327
 328     :param str name: name of the field
 329     :param str endianness: format endianness string; one of "<", ">", or "="
 330     :param str format: a single format character
 331     """
 332
 333     __slots__ = ["packer"]
 334     def __init__(self, name, endianity, format):
 335         if endianity not in (">", "<", "="):
 336             raise ValueError("endianity must be be '=', '<', or '>'",
 337                 endianity)
 338         if len(format) != 1:
 339             raise ValueError("must specify one and only one format char")
 340         self.packer = Packer(endianity + format)
 341         StaticField.__init__(self, name, self.packer.size)
 342     def __getstate__(self):
 343         attrs = StaticField.__getstate__(self)
 344         attrs["packer"] = attrs["packer"].format
 345         return attrs
 346     def __setstate__(self, attrs):
 347         attrs["packer"] = Packer(attrs["packer"])
 348         return StaticField.__setstate__(attrs)
 349     def _parse(self, stream, context):
 350         try:
 351             return self.packer.unpack(_read_stream(stream, self.length))[0]
 352         except Exception as ex:
 353             raise FieldError(ex)
 354     def _build(self, obj, stream, context):
 355         try:
 356             _write_stream(stream, self.length, self.packer.pack(obj))
 357         except Exception as ex:
 358             raise FieldError(ex)
 359
 360 class MetaField(Construct):
 361     """
 362     A variable-length field. The length is obtained at runtime from a
 363     function.
 364
 365     :param str name: name of the field
 366     :param callable lengthfunc: callable that takes a context and returns
 367                                 length as an int
 368
 369     >>> foo = Struct("foo",
 370     ...     Byte("length"),
 371     ...     MetaField("data", lambda ctx: ctx["length"])
 372     ... )
 373     >>> foo.parse("\\x03ABC")
 374     Container(data = 'ABC', length = 3)
 375     >>> foo.parse("\\x04ABCD")
 376     Container(data = 'ABCD', length = 4)
 377     """
 378
 379     __slots__ = ["lengthfunc"]
 380     def __init__(self, name, lengthfunc):
 381         Construct.__init__(self, name)
 382         self.lengthfunc = lengthfunc
 383         self._set_flag(self.FLAG_DYNAMIC)
 384     def _parse(self, stream, context):
 385         return _read_stream(stream, self.lengthfunc(context))
 386     def _build(self, obj, stream, context):
 387         _write_stream(stream, self.lengthfunc(context), obj)
 388     def _sizeof(self, context):
 389         return self.lengthfunc(context)
 390
 391
 392 #===============================================================================
 393 # arrays and repeaters
 394 #===============================================================================
 395 class MetaArray(Subconstruct):
 396     """
 397     An array (repeater) of a meta-count. The array will iterate exactly
 398     `countfunc()` times. Will raise ArrayError if less elements are found.
 399     See also Array, Range and RepeatUntil.
 400
 401     Parameters:
 402     * countfunc - a function that takes the context as a parameter and returns
 403       the number of elements of the array (count)
 404     * subcon - the subcon to repeat `countfunc()` times
 405
 406     Example:
 407     MetaArray(lambda ctx: 5, UBInt8("foo"))
 408     """
 409     __slots__ = ["countfunc"]
 410     def __init__(self, countfunc, subcon):
 411         Subconstruct.__init__(self, subcon)
 412         self.countfunc = countfunc
 413         self._clear_flag(self.FLAG_COPY_CONTEXT)
 414         self._set_flag(self.FLAG_DYNAMIC)
 415     def _parse(self, stream, context):
 416         obj = ListContainer()
 417         c = 0
 418         count = self.countfunc(context)
 419         try:
 420             if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 421                 while c < count:
 422                     obj.append(self.subcon._parse(stream, context.__copy__()))
 423                     c += 1
 424             else:
 425                 while c < count:
 426                     obj.append(self.subcon._parse(stream, context))
 427                     c += 1
 428         except ConstructError as ex:
 429             raise ArrayError("expected %d, found %d" % (count, c), ex)
 430         return obj
 431     def _build(self, obj, stream, context):
 432         count = self.countfunc(context)
 433         if len(obj) != count:
 434             raise ArrayError("expected %d, found %d" % (count, len(obj)))
 435         if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 436             for subobj in obj:
 437                 self.subcon._build(subobj, stream, context.__copy__())
 438         else:
 439             for subobj in obj:
 440                 self.subcon._build(subobj, stream, context)
 441     def _sizeof(self, context):
 442         return self.subcon._sizeof(context) * self.countfunc(context)
 443
 444 class Range(Subconstruct):
 445     """
 446     A range-array. The subcon will iterate between `mincount` to `maxcount`
 447     times. If less than `mincount` elements are found, raises RangeError.
 448     See also GreedyRange and OptionalGreedyRange.
 449
 450     The general-case repeater. Repeats the given unit for at least mincount
 451     times, and up to maxcount times. If an exception occurs (EOF, validation
 452     error), the repeater exits. If less than mincount units have been
 453     successfully parsed, a RangeError is raised.
 454
 455     .. note::
 456        This object requires a seekable stream for parsing.
 457
 458     :param int mincount: the minimal count
 459     :param int maxcount: the maximal count
 460     :param Construct subcon: the subcon to repeat
 461
 462     >>> c = Range(3, 7, UBInt8("foo"))
 463     >>> c.parse("\\x01\\x02")
 464     Traceback (most recent call last):
 465       ...
 466     construct.core.RangeError: expected 3..7, found 2
 467     >>> c.parse("\\x01\\x02\\x03")
 468     [1, 2, 3]
 469     >>> c.parse("\\x01\\x02\\x03\\x04\\x05\\x06")
 470     [1, 2, 3, 4, 5, 6]
 471     >>> c.parse("\\x01\\x02\\x03\\x04\\x05\\x06\\x07")
 472     [1, 2, 3, 4, 5, 6, 7]
 473     >>> c.parse("\\x01\\x02\\x03\\x04\\x05\\x06\\x07\\x08\\x09")
 474     [1, 2, 3, 4, 5, 6, 7]
 475     >>> c.build([1,2])
 476     Traceback (most recent call last):
 477       ...
 478     construct.core.RangeError: expected 3..7, found 2
 479     >>> c.build([1,2,3,4])
 480     '\\x01\\x02\\x03\\x04'
 481     >>> c.build([1,2,3,4,5,6,7,8])
 482     Traceback (most recent call last):
 483       ...
 484     construct.core.RangeError: expected 3..7, found 8
 485     """
 486
 487     __slots__ = ["mincount", "maxcout"]
 488     def __init__(self, mincount, maxcout, subcon):
 489         Subconstruct.__init__(self, subcon)
 490         self.mincount = mincount
 491         self.maxcout = maxcout
 492         self._clear_flag(self.FLAG_COPY_CONTEXT)
 493         self._set_flag(self.FLAG_DYNAMIC)
 494     def _parse(self, stream, context):
 495         obj = ListContainer()
 496         c = 0
 497         try:
 498             if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 499                 while c < self.maxcout:
 500                     pos = stream.tell()
 501                     obj.append(self.subcon._parse(stream, context.__copy__()))
 502                     c += 1
 503             else:
 504                 while c < self.maxcout:
 505                     pos = stream.tell()
 506                     obj.append(self.subcon._parse(stream, context))
 507                     c += 1
 508         except ConstructError as ex:
 509             if c < self.mincount:
 510                 raise RangeError("expected %d to %d, found %d" %
 511                     (self.mincount, self.maxcout, c), ex)
 512             stream.seek(pos)
 513         return obj
 514     def _build(self, obj, stream, context):
 515         if len(obj) < self.mincount or len(obj) > self.maxcout:
 516             raise RangeError("expected %d to %d, found %d" %
 517                 (self.mincount, self.maxcout, len(obj)))
 518         cnt = 0
 519         try:
 520             if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 521                 for subobj in obj:
 522                     if isinstance(obj, bytes):
 523                         subobj = bchr(subobj)
 524                     self.subcon._build(subobj, stream, context.__copy__())
 525                     cnt += 1
 526             else:
 527                 for subobj in obj:
 528                     if isinstance(obj, bytes):
 529                         subobj = bchr(subobj)
 530                     self.subcon._build(subobj, stream, context)
 531                     cnt += 1
 532         except ConstructError as ex:
 533             if cnt < self.mincount:
 534                 raise RangeError("expected %d to %d, found %d" %
 535                     (self.mincount, self.maxcout, len(obj)), ex)
 536     def _sizeof(self, context):
 537         raise SizeofError("can't calculate size")
 538
 539 class RepeatUntil(Subconstruct):
 540     """
 541     An array that repeats until the predicate indicates it to stop. Note that
 542     the last element (which caused the repeat to exit) is included in the
 543     return value.
 544
 545     Parameters:
 546     * predicate - a predicate function that takes (obj, context) and returns
 547       True if the stop-condition is met, or False to continue.
 548     * subcon - the subcon to repeat.
 549
 550     Example:
 551     # will read chars until b\x00 (inclusive)
 552     RepeatUntil(lambda obj, ctx: obj == b"\x00",
 553         Field("chars", 1)
 554     )
 555     """
 556     __slots__ = ["predicate"]
 557     def __init__(self, predicate, subcon):
 558         Subconstruct.__init__(self, subcon)
 559         self.predicate = predicate
 560         self._clear_flag(self.FLAG_COPY_CONTEXT)
 561         self._set_flag(self.FLAG_DYNAMIC)
 562     def _parse(self, stream, context):
 563         obj = []
 564         try:
 565             if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 566                 while True:
 567                     subobj = self.subcon._parse(stream, context.__copy__())
 568                     obj.append(subobj)
 569                     if self.predicate(subobj, context):
 570                         break
 571             else:
 572                 while True:
 573                     subobj = self.subcon._parse(stream, context)
 574                     obj.append(subobj)
 575                     if self.predicate(subobj, context):
 576                         break
 577         except ConstructError as ex:
 578             raise ArrayError("missing terminator", ex)
 579         return obj
 580     def _build(self, obj, stream, context):
 581         terminated = False
 582         if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
 583             for subobj in obj:
 584                 self.subcon._build(subobj, stream, context.__copy__())
 585                 if self.predicate(subobj, context):
 586                     terminated = True
 587                     break
 588         else:
 589             for subobj in obj:
 590                 subobj = bchr(subobj)
 591                 self.subcon._build(subobj, stream, context.__copy__())
 592                 if self.predicate(subobj, context):
 593                     terminated = True
 594                     break
 595         if not terminated:
 596             raise ArrayError("missing terminator")
 597     def _sizeof(self, context):
 598         raise SizeofError("can't calculate size")
 599
 600
 601 #===============================================================================
 602 # structures and sequences
 603 #===============================================================================
 604 class Struct(Construct):
 605     """
 606     A sequence of named constructs, similar to structs in C. The elements are
 607     parsed and built in the order they are defined.
 608     See also Embedded.
 609
 610     Parameters:
 611     * name - the name of the structure
 612     * subcons - a sequence of subconstructs that make up this structure.
 613     * nested - a keyword-only argument that indicates whether this struct
 614       creates a nested context. The default is True. This parameter is
 615       considered "advanced usage", and may be removed in the future.
 616
 617     Example:
 618     Struct("foo",
 619         UBInt8("first_element"),
 620         UBInt16("second_element"),
 621         Padding(2),
 622         UBInt8("third_element"),
 623     )
 624     """
 625     __slots__ = ["subcons", "nested"]
 626     def __init__(self, name, *subcons, **kw):
 627         self.nested = kw.pop("nested", True)
 628         if kw:
 629             raise TypeError("the only keyword argument accepted is 'nested'", kw)
 630         Construct.__init__(self, name)
 631         self.subcons = subcons
 632         self._inherit_flags(*subcons)
 633         self._clear_flag(self.FLAG_EMBED)
 634     def _parse(self, stream, context):
 635         if "<obj>" in context:
 636             obj = context["<obj>"]
 637             del context["<obj>"]
 638         else:
 639             obj = Container()
 640             if self.nested:
 641                 context = Container(_ = context)
 642         for sc in self.subcons:
 643             if sc.conflags & self.FLAG_EMBED:
 644                 context["<obj>"] = obj
 645                 sc._parse(stream, context)
 646             else:
 647                 subobj = sc._parse(stream, context)
 648                 if sc.name is not None:
 649                     obj[sc.name] = subobj
 650                     context[sc.name] = subobj
 651         return obj
 652     def _build(self, obj, stream, context):
 653         if "<unnested>" in context:
 654             del context["<unnested>"]
 655         elif self.nested:
 656             context = Container(_ = context)
 657         for sc in self.subcons:
 658             if sc.conflags & self.FLAG_EMBED:
 659                 context["<unnested>"] = True
 660                 subobj = obj
 661             elif sc.name is None:
 662                 subobj = None
 663             else:
 664                 subobj = getattr(obj, sc.name)
 665                 context[sc.name] = subobj
 666             sc._build(subobj, stream, context)
 667     def _sizeof(self, context):
 668         if self.nested:
 669             context = Container(_ = context)
 670         return sum(sc._sizeof(context) for sc in self.subcons)
 671
 672 class Sequence(Struct):
 673     """
 674     A sequence of unnamed constructs. The elements are parsed and built in the
 675     order they are defined.
 676     See also Embedded.
 677
 678     Parameters:
 679     * name - the name of the structure
 680     * subcons - a sequence of subconstructs that make up this structure.
 681     * nested - a keyword-only argument that indicates whether this struct
 682       creates a nested context. The default is True. This parameter is
 683       considered "advanced usage", and may be removed in the future.
 684
 685     Example:
 686     Sequence("foo",
 687         UBInt8("first_element"),
 688         UBInt16("second_element"),
 689         Padding(2),
 690         UBInt8("third_element"),
 691     )
 692     """
 693     __slots__ = []
 694     def _parse(self, stream, context):
 695         if "<obj>" in context:
 696             obj = context["<obj>"]
 697             del context["<obj>"]
 698         else:
 699             obj = ListContainer()
 700             if self.nested:
 701                 context = Container(_ = context)
 702         for sc in self.subcons:
 703             if sc.conflags & self.FLAG_EMBED:
 704                 context["<obj>"] = obj
 705                 sc._parse(stream, context)
 706             else:
 707                 subobj = sc._parse(stream, context)
 708                 if sc.name is not None:
 709                     obj.append(subobj)
 710                     context[sc.name] = subobj
 711         return obj
 712     def _build(self, obj, stream, context):
 713         if "<unnested>" in context:
 714             del context["<unnested>"]
 715         elif self.nested:
 716             context = Container(_ = context)
 717         objiter = iter(obj)
 718         for sc in self.subcons:
 719             if sc.conflags & self.FLAG_EMBED:
 720                 context["<unnested>"] = True
 721                 subobj = objiter
 722             elif sc.name is None:
 723                 subobj = None
 724             else:
 725                 subobj = advance_iterator(objiter)
 726                 context[sc.name] = subobj
 727             sc._build(subobj, stream, context)
 728
 729 class Union(Construct):
 730     """
 731     a set of overlapping fields (like unions in C). when parsing,
 732     all fields read the same data; when building, only the first subcon
 733     (called "master") is used.
 734
 735     Parameters:
 736     * name - the name of the union
 737     * master - the master subcon, i.e., the subcon used for building and
 738       calculating the total size
 739     * subcons - additional subcons
 740
 741     Example:
 742     Union("what_are_four_bytes",
 743         UBInt32("one_dword"),
 744         Struct("two_words", UBInt16("first"), UBInt16("second")),
 745         Struct("four_bytes",
 746             UBInt8("a"),
 747             UBInt8("b"),
 748             UBInt8("c"),
 749             UBInt8("d")
 750         ),
 751     )
 752     """
 753     __slots__ = ["parser", "builder"]
 754     def __init__(self, name, master, *subcons, **kw):
 755         Construct.__init__(self, name)
 756         args = [Peek(sc) for sc in subcons]
 757         args.append(MetaField(None, lambda ctx: master._sizeof(ctx)))
 758         self.parser = Struct(name, Peek(master, perform_build = True), *args)
 759         self.builder = Struct(name, master)
 760     def _parse(self, stream, context):
 761         return self.parser._parse(stream, context)
 762     def _build(self, obj, stream, context):
 763         return self.builder._build(obj, stream, context)
 764     def _sizeof(self, context):
 765         return self.builder._sizeof(context)
 766
 767 #===============================================================================
 768 # conditional
 769 #===============================================================================
 770 class Switch(Construct):
 771     """
 772     A conditional branch. Switch will choose the case to follow based on
 773     the return value of keyfunc. If no case is matched, and no default value
 774     is given, SwitchError will be raised.
 775     See also Pass.
 776
 777     Parameters:
 778     * name - the name of the construct
 779     * keyfunc - a function that takes the context and returns a key, which
 780       will ne used to choose the relevant case.
 781     * cases - a dictionary mapping keys to constructs. the keys can be any
 782       values that may be returned by keyfunc.
 783     * default - a default value to use when the key is not found in the cases.
 784       if not supplied, an exception will be raised when the key is not found.
 785       You can use the builtin construct Pass for 'do-nothing'.
 786     * include_key - whether or not to include the key in the return value
 787       of parsing. defualt is False.
 788
 789     Example:
 790     Struct("foo",
 791         UBInt8("type"),
 792         Switch("value", lambda ctx: ctx.type, {
 793                 1 : UBInt8("spam"),
 794                 2 : UBInt16("spam"),
 795                 3 : UBInt32("spam"),
 796                 4 : UBInt64("spam"),
 797             }
 798         ),
 799     )
 800     """
 801
 802     class NoDefault(Construct):
 803         def _parse(self, stream, context):
 804             raise SwitchError("no default case defined")
 805         def _build(self, obj, stream, context):
 806             raise SwitchError("no default case defined")
 807         def _sizeof(self, context):
 808             raise SwitchError("no default case defined")
 809     NoDefault = NoDefault("No default value specified")
 810
 811     __slots__ = ["subcons", "keyfunc", "cases", "default", "include_key"]
 812
 813     def __init__(self, name, keyfunc, cases, default = NoDefault,
 814                  include_key = False):
 815         Construct.__init__(self, name)
 816         self._inherit_flags(*cases.values())
 817         self.keyfunc = keyfunc
 818         self.cases = cases
 819         self.default = default
 820         self.include_key = include_key
 821         self._inherit_flags(*cases.values())
 822         self._set_flag(self.FLAG_DYNAMIC)
 823     def _parse(self, stream, context):
 824         key = self.keyfunc(context)
 825         obj = self.cases.get(key, self.default)._parse(stream, context)
 826         if self.include_key:
 827             return key, obj
 828         else:
 829             return obj
 830     def _build(self, obj, stream, context):
 831         if self.include_key:
 832             key, obj = obj
 833         else:
 834             key = self.keyfunc(context)
 835         case = self.cases.get(key, self.default)
 836         case._build(obj, stream, context)
 837     def _sizeof(self, context):
 838         case = self.cases.get(self.keyfunc(context), self.default)
 839         return case._sizeof(context)
 840
 841 class Select(Construct):
 842     """
 843     Selects the first matching subconstruct. It will literally try each of
 844     the subconstructs, until one matches.
 845
 846     Notes:
 847     * requires a seekable stream.
 848
 849     Parameters:
 850     * name - the name of the construct
 851     * subcons - the subcons to try (order-sensitive)
 852     * include_name - a keyword only argument, indicating whether to include
 853       the name of the selected subcon in the return value of parsing. default
 854       is false.
 855
 856     Example:
 857     Select("foo",
 858         UBInt64("large"),
 859         UBInt32("medium"),
 860         UBInt16("small"),
 861         UBInt8("tiny"),
 862     )
 863     """
 864     __slots__ = ["subcons", "include_name"]
 865     def __init__(self, name, *subcons, **kw):
 866         include_name = kw.pop("include_name", False)
 867         if kw:
 868             raise TypeError("the only keyword argument accepted "
 869                 "is 'include_name'", kw)
 870         Construct.__init__(self, name)
 871         self.subcons = subcons
 872         self.include_name = include_name
 873         self._inherit_flags(*subcons)
 874         self._set_flag(self.FLAG_DYNAMIC)
 875     def _parse(self, stream, context):
 876         for sc in self.subcons:
 877             pos = stream.tell()
 878             context2 = context.__copy__()
 879             try:
 880                 obj = sc._parse(stream, context2)
 881             except ConstructError:
 882                 stream.seek(pos)
 883             else:
 884                 context.__update__(context2)
 885                 if self.include_name:
 886                     return sc.name, obj
 887                 else:
 888                     return obj
 889         raise SelectError("no subconstruct matched")
 890     def _build(self, obj, stream, context):
 891         if self.include_name:
 892             name, obj = obj
 893             for sc in self.subcons:
 894                 if sc.name == name:
 895                     sc._build(obj, stream, context)
 896                     return
 897         else:
 898             for sc in self.subcons:
 899                 stream2 = BytesIO()
 900                 context2 = context.__copy__()
 901                 try:
 902                     sc._build(obj, stream2, context2)
 903                 except Exception:
 904                     pass
 905                 else:
 906                     context.__update__(context2)
 907                     stream.write(stream2.getvalue())
 908                     return
 909         raise SelectError("no subconstruct matched", obj)
 910     def _sizeof(self, context):
 911         raise SizeofError("can't calculate size")
 912
 913
 914 #===============================================================================
 915 # stream manipulation
 916 #===============================================================================
 917 class Pointer(Subconstruct):
 918     """
 919     Changes the stream position to a given offset, where the construction
 920     should take place, and restores the stream position when finished.
 921     See also Anchor, OnDemand and OnDemandPointer.
 922
 923     Notes:
 924     * requires a seekable stream.
 925
 926     Parameters:
 927     * offsetfunc: a function that takes the context and returns an absolute
 928       stream position, where the construction would take place
 929     * subcon - the subcon to use at `offsetfunc()`
 930
 931     Example:
 932     Struct("foo",
 933         UBInt32("spam_pointer"),
 934         Pointer(lambda ctx: ctx.spam_pointer,
 935             Array(5, UBInt8("spam"))
 936         )
 937     )
 938     """
 939     __slots__ = ["offsetfunc"]
 940     def __init__(self, offsetfunc, subcon):
 941         Subconstruct.__init__(self, subcon)
 942         self.offsetfunc = offsetfunc
 943     def _parse(self, stream, context):
 944         newpos = self.offsetfunc(context)
 945         origpos = stream.tell()
 946         stream.seek(newpos)
 947         obj = self.subcon._parse(stream, context)
 948         stream.seek(origpos)
 949         return obj
 950     def _build(self, obj, stream, context):
 951         newpos = self.offsetfunc(context)
 952         origpos = stream.tell()
 953         stream.seek(newpos)
 954         self.subcon._build(obj, stream, context)
 955         stream.seek(origpos)
 956     def _sizeof(self, context):
 957         return 0
 958
 959 class Peek(Subconstruct):
 960     """
 961     Peeks at the stream: parses without changing the stream position.
 962     See also Union. If the end of the stream is reached when peeking,
 963     returns None.
 964
 965     Notes:
 966     * requires a seekable stream.
 967
 968     Parameters:
 969     * subcon - the subcon to peek at
 970     * perform_build - whether or not to perform building. by default this
 971       parameter is set to False, meaning building is a no-op.
 972
 973     Example:
 974     Peek(UBInt8("foo"))
 975     """
 976     __slots__ = ["perform_build"]
 977     def __init__(self, subcon, perform_build = False):
 978         Subconstruct.__init__(self, subcon)
 979         self.perform_build = perform_build
 980     def _parse(self, stream, context):
 981         pos = stream.tell()
 982         try:
 983             return self.subcon._parse(stream, context)
 984         except FieldError:
 985             pass
 986         finally:
 987             stream.seek(pos)
 988     def _build(self, obj, stream, context):
 989         if self.perform_build:
 990             self.subcon._build(obj, stream, context)
 991     def _sizeof(self, context):
 992         return 0
 993
 994 class OnDemand(Subconstruct):
 995     """
 996     Allows for on-demand (lazy) parsing. When parsing, it will return a
 997     LazyContainer that represents a pointer to the data, but does not actually
 998     parses it from stream until it's "demanded".
 999     By accessing the 'value' property of LazyContainers, you will demand the
1000     data from the stream. The data will be parsed and cached for later use.
1001     You can use the 'has_value' property to know whether the data has already
1002     been demanded.
1003     See also OnDemandPointer.
1004
1005     Notes:
1006     * requires a seekable stream.
1007
1008     Parameters:
1009     * subcon -
1010     * advance_stream - whether or not to advance the stream position. by
1011       default this is True, but if subcon is a pointer, this should be False.
1012     * force_build - whether or not to force build. If set to False, and the
1013       LazyContainer has not been demaned, building is a no-op.
1014
1015     Example:
1016     OnDemand(Array(10000, UBInt8("foo"))
1017     """
1018     __slots__ = ["advance_stream", "force_build"]
1019     def __init__(self, subcon, advance_stream = True, force_build = True):
1020         Subconstruct.__init__(self, subcon)
1021         self.advance_stream = advance_stream
1022         self.force_build = force_build
1023     def _parse(self, stream, context):
1024         obj = LazyContainer(self.subcon, stream, stream.tell(), context)
1025         if self.advance_stream:
1026             stream.seek(self.subcon._sizeof(context), 1)
1027         return obj
1028     def _build(self, obj, stream, context):
1029         if not isinstance(obj, LazyContainer):
1030             self.subcon._build(obj, stream, context)
1031         elif self.force_build or obj.has_value:
1032             self.subcon._build(obj.value, stream, context)
1033         elif self.advance_stream:
1034             stream.seek(self.subcon._sizeof(context), 1)
1035
1036 class Buffered(Subconstruct):
1037     """
1038     Creates an in-memory buffered stream, which can undergo encoding and
1039     decoding prior to being passed on to the subconstruct.
1040     See also Bitwise.
1041
1042     Note:
1043     * Do not use pointers inside Buffered
1044
1045     Parameters:
1046     * subcon - the subcon which will operate on the buffer
1047     * encoder - a function that takes a string and returns an encoded
1048       string (used after building)
1049     * decoder - a function that takes a string and returns a decoded
1050       string (used before parsing)
1051     * resizer - a function that takes the size of the subcon and "adjusts"
1052       or "resizes" it according to the encoding/decoding process.
1053
1054     Example:
1055     Buffered(BitField("foo", 16),
1056         encoder = decode_bin,
1057         decoder = encode_bin,
1058         resizer = lambda size: size / 8,
1059     )
1060     """
1061     __slots__ = ["encoder", "decoder", "resizer"]
1062     def __init__(self, subcon, decoder, encoder, resizer):
1063         Subconstruct.__init__(self, subcon)
1064         self.encoder = encoder
1065         self.decoder = decoder
1066         self.resizer = resizer
1067     def _parse(self, stream, context):
1068         data = _read_stream(stream, self._sizeof(context))
1069         stream2 = BytesIO(self.decoder(data))
1070         return self.subcon._parse(stream2, context)
1071     def _build(self, obj, stream, context):
1072         size = self._sizeof(context)
1073         stream2 = BytesIO()
1074         self.subcon._build(obj, stream2, context)
1075         data = self.encoder(stream2.getvalue())
1076         assert len(data) == size
1077         _write_stream(stream, self._sizeof(context), data)
1078     def _sizeof(self, context):
1079         return self.resizer(self.subcon._sizeof(context))
1080
1081 class Restream(Subconstruct):
1082     """
1083     Wraps the stream with a read-wrapper (for parsing) or a
1084     write-wrapper (for building). The stream wrapper can buffer the data
1085     internally, reading it from- or writing it to the underlying stream
1086     as needed. For example, BitStreamReader reads whole bytes from the
1087     underlying stream, but returns them as individual bits.
1088     See also Bitwise.
1089
1090     When the parsing or building is done, the stream's close method
1091     will be invoked. It can perform any finalization needed for the stream
1092     wrapper, but it must not close the underlying stream.
1093
1094     Note:
1095     * Do not use pointers inside Restream
1096
1097     Parameters:
1098     * subcon - the subcon
1099     * stream_reader - the read-wrapper
1100     * stream_writer - the write wrapper
1101     * resizer - a function that takes the size of the subcon and "adjusts"
1102       or "resizes" it according to the encoding/decoding process.
1103
1104     Example:
1105     Restream(BitField("foo", 16),
1106         stream_reader = BitStreamReader,
1107         stream_writer = BitStreamWriter,
1108         resizer = lambda size: size / 8,
1109     )
1110     """
1111     __slots__ = ["stream_reader", "stream_writer", "resizer"]
1112     def __init__(self, subcon, stream_reader, stream_writer, resizer):
1113         Subconstruct.__init__(self, subcon)
1114         self.stream_reader = stream_reader
1115         self.stream_writer = stream_writer
1116         self.resizer = resizer
1117     def _parse(self, stream, context):
1118         stream2 = self.stream_reader(stream)
1119         obj = self.subcon._parse(stream2, context)
1120         stream2.close()
1121         return obj
1122     def _build(self, obj, stream, context):
1123         stream2 = self.stream_writer(stream)
1124         self.subcon._build(obj, stream2, context)
1125         stream2.close()
1126     def _sizeof(self, context):
1127         return self.resizer(self.subcon._sizeof(context))
1128
1129
1130 #===============================================================================
1131 # miscellaneous
1132 #===============================================================================
1133 class Reconfig(Subconstruct):
1134     """
1135     Reconfigures a subconstruct. Reconfig can be used to change the name and
1136     set and clear flags of the inner subcon.
1137
1138     Parameters:
1139     * name - the new name
1140     * subcon - the subcon to reconfigure
1141     * setflags - the flags to set (default is 0)
1142     * clearflags - the flags to clear (default is 0)
1143
1144     Example:
1145     Reconfig("foo", UBInt8("bar"))
1146     """
1147     __slots__ = []
1148     def __init__(self, name, subcon, setflags = 0, clearflags = 0):
1149         Construct.__init__(self, name, subcon.conflags)
1150         self.subcon = subcon
1151         self._set_flag(setflags)
1152         self._clear_flag(clearflags)
1153
1154 class Anchor(Construct):
1155     """
1156     Returns the "anchor" (stream position) at the point where it's inserted.
1157     Useful for adjusting relative offsets to absolute positions, or to measure
1158     sizes of constructs.
1159     absolute pointer = anchor + relative offset
1160     size = anchor_after - anchor_before
1161     See also Pointer.
1162
1163     Notes:
1164     * requires a seekable stream.
1165
1166     Parameters:
1167     * name - the name of the anchor
1168
1169     Example:
1170     Struct("foo",
1171         Anchor("base"),
1172         UBInt8("relative_offset"),
1173         Pointer(lambda ctx: ctx.relative_offset + ctx.base,
1174             UBInt8("data")
1175         )
1176     )
1177     """
1178     __slots__ = []
1179     def _parse(self, stream, context):
1180         return stream.tell()
1181     def _build(self, obj, stream, context):
1182         context[self.name] = stream.tell()
1183     def _sizeof(self, context):
1184         return 0
1185
1186 class Value(Construct):
1187     """
1188     A computed value.
1189
1190     Parameters:
1191     * name - the name of the value
1192     * func - a function that takes the context and return the computed value
1193
1194     Example:
1195     Struct("foo",
1196         UBInt8("width"),
1197         UBInt8("height"),
1198         Value("total_pixels", lambda ctx: ctx.width * ctx.height),
1199     )
1200     """
1201     __slots__ = ["func"]
1202     def __init__(self, name, func):
1203         Construct.__init__(self, name)
1204         self.func = func
1205         self._set_flag(self.FLAG_DYNAMIC)
1206     def _parse(self, stream, context):
1207         return self.func(context)
1208     def _build(self, obj, stream, context):
1209         context[self.name] = self.func(context)
1210     def _sizeof(self, context):
1211         return 0
1212
1213 #class Dynamic(Construct):
1214 #    """
1215 #    Dynamically creates a construct and uses it for parsing and building.
1216 #    This allows you to create change the construction tree on the fly.
1217 #    Deprecated.
1218 #
1219 #    Parameters:
1220 #    * name - the name of the construct
1221 #    * factoryfunc - a function that takes the context and returns a new
1222 #      construct object which will be used for parsing and building.
1223 #
1224 #    Example:
1225 #    def factory(ctx):
1226 #        if ctx.bar == 8:
1227 #            return UBInt8("spam")
1228 #        if ctx.bar == 9:
1229 #            return String("spam", 9)
1230 #
1231 #    Struct("foo",
1232 #        UBInt8("bar"),
1233 #        Dynamic("spam", factory),
1234 #    )
1235 #    """
1236 #    __slots__ = ["factoryfunc"]
1237 #    def __init__(self, name, factoryfunc):
1238 #        Construct.__init__(self, name, self.FLAG_COPY_CONTEXT)
1239 #        self.factoryfunc = factoryfunc
1240 #        self._set_flag(self.FLAG_DYNAMIC)
1241 #    def _parse(self, stream, context):
1242 #        return self.factoryfunc(context)._parse(stream, context)
1243 #    def _build(self, obj, stream, context):
1244 #        return self.factoryfunc(context)._build(obj, stream, context)
1245 #    def _sizeof(self, context):
1246 #        return self.factoryfunc(context)._sizeof(context)
1247
1248 class LazyBound(Construct):
1249     """
1250     Lazily bound construct, useful for constructs that need to make cyclic
1251     references (linked-lists, expression trees, etc.).
1252
1253     Parameters:
1254
1255
1256     Example:
1257     foo = Struct("foo",
1258         UBInt8("bar"),
1259         LazyBound("next", lambda: foo),
1260     )
1261     """
1262     __slots__ = ["bindfunc", "bound"]
1263     def __init__(self, name, bindfunc):
1264         Construct.__init__(self, name)
1265         self.bound = None
1266         self.bindfunc = bindfunc
1267     def _parse(self, stream, context):
1268         if self.bound is None:
1269             self.bound = self.bindfunc()
1270         return self.bound._parse(stream, context)
1271     def _build(self, obj, stream, context):
1272         if self.bound is None:
1273             self.bound = self.bindfunc()
1274         self.bound._build(obj, stream, context)
1275     def _sizeof(self, context):
1276         if self.bound is None:
1277             self.bound = self.bindfunc()
1278         return self.bound._sizeof(context)
1279
1280 class Pass(Construct):
1281     """
1282     A do-nothing construct, useful as the default case for Switch, or
1283     to indicate Enums.
1284     See also Switch and Enum.
1285
1286     Notes:
1287     * this construct is a singleton. do not try to instatiate it, as it
1288       will not work...
1289
1290     Example:
1291     Pass
1292     """
1293     __slots__ = []
1294     def _parse(self, stream, context):
1295         pass
1296     def _build(self, obj, stream, context):
1297         assert obj is None
1298     def _sizeof(self, context):
1299         return 0
1300 Pass = Pass(None)
1301
1302 class Terminator(Construct):
1303     """
1304     Asserts the end of the stream has been reached at the point it's placed.
1305     You can use this to ensure no more unparsed data follows.
1306
1307     Notes:
1308     * this construct is only meaningful for parsing. for building, it's
1309       a no-op.
1310     * this construct is a singleton. do not try to instatiate it, as it
1311       will not work...
1312
1313     Example:
1314     Terminator
1315     """
1316     __slots__ = []
1317     def _parse(self, stream, context):
1318         if stream.read(1):
1319             raise TerminatorError("expected end of stream")
1320     def _build(self, obj, stream, context):
1321         assert obj is None
1322     def _sizeof(self, context):
1323         return 0
1324 Terminator = Terminator(None)