third_party/pigweed/repo/pw_tokenizer/py/pw_tokenizer/decode.py

   1 # Copyright 2020 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """Decodes arguments and formats tokenized messages.
  15
  16 The decode(format_string, encoded_arguments) function provides a simple way to
  17 format a string with encoded arguments. The FormatString class may also be used.
  18
  19 Missing, truncated, or otherwise corrupted arguments are handled and displayed
  20 in the resulting string with an error message.
  21 """
  22
  23 import re
  24 import struct
  25 from typing import Iterable, List, NamedTuple, Match, Sequence, Tuple
  26
  27
  28 def zigzag_decode(value: int) -> int:
  29     """ZigZag decode function from protobuf's wire_format module."""
  30     if not value & 0x1:
  31         return value >> 1
  32     return (value >> 1) ^ (~0)
  33
  34
  35 class FormatSpec:
  36     """Represents a format specifier parsed from a printf-style string."""
  37
  38     # Regular expression for finding format specifiers.
  39     FORMAT_SPEC = re.compile(r'%(?:(?P<flags>[+\- #0]*\d*(?:\.\d+)?)'
  40                              r'(?P<length>hh|h|ll|l|j|z|t|L)?'
  41                              r'(?P<type>[csdioxXufFeEaAgGnp])|%)')
  42
  43     # Conversions to make format strings Python compatible.
  44     _UNSUPPORTED_LENGTH = frozenset(['hh', 'll', 'j', 'z', 't'])
  45     _REMAP_TYPE = {'a': 'f', 'A': 'F'}
  46
  47     # Conversion specifiers by type; n is not supported.
  48     _SIGNED_INT = 'di'
  49     _UNSIGNED_INT = frozenset('oxXup')
  50     _FLOATING_POINT = frozenset('fFeEaAgG')
  51
  52     _PACKED_FLOAT = struct.Struct('<f')
  53
  54     @classmethod
  55     def from_string(cls, format_specifier: str):
  56         """Creates a FormatSpec from a str with a single format specifier."""
  57         match = cls.FORMAT_SPEC.fullmatch(format_specifier)
  58
  59         if not match:
  60             raise ValueError(
  61                 '{!r} is not a valid single format specifier'.format(
  62                     format_specifier))
  63
  64         return cls(match)
  65
  66     def __init__(self, re_match: Match):
  67         """Constructs a FormatSpec from an re.Match object for FORMAT_SPEC."""
  68         self.match = re_match
  69         self.specifier: str = self.match.group()
  70
  71         self.flags: str = self.match.group('flags') or ''
  72         self.length: str = self.match.group('length') or ''
  73
  74         # If there is no type, the format spec is %%.
  75         self.type: str = self.match.group('type') or '%'
  76
  77         # %p prints as 0xFEEDBEEF; other specs may need length/type switched
  78         if self.type == 'p':
  79             self.compatible = '0x%08X'
  80         else:
  81             self.compatible = ''.join([
  82                 '%', self.flags,
  83                 '' if self.length in self._UNSUPPORTED_LENGTH else '',
  84                 self._REMAP_TYPE.get(self.type, self.type)
  85             ])
  86
  87     def decode(self, encoded_arg: bytes) -> 'DecodedArg':
  88         """Decodes the provided data according to this format specifier."""
  89         if self.type == '%':  # literal %
  90             return DecodedArg(self, (),
  91                               b'')  # Use () as the value for % formatting.
  92
  93         if self.type == 's':  # string
  94             return self._decode_string(encoded_arg)
  95
  96         if self.type == 'c':  # character
  97             return self._decode_char(encoded_arg)
  98
  99         if self.type in self._SIGNED_INT:
 100             return self._decode_signed_integer(encoded_arg)
 101
 102         if self.type in self._UNSIGNED_INT:
 103             return self._decode_unsigned_integer(encoded_arg)
 104
 105         if self.type in self._FLOATING_POINT:
 106             return self._decode_float(encoded_arg)
 107
 108         # Unsupported specifier (e.g. %n)
 109         return DecodedArg(
 110             self, None, b'', DecodedArg.DECODE_ERROR,
 111             'Unsupported conversion specifier "{}"'.format(self.type))
 112
 113     def _decode_signed_integer(self, encoded: bytes) -> 'DecodedArg':
 114         """Decodes a signed variable-length integer."""
 115         if not encoded:
 116             return DecodedArg.missing(self)
 117
 118         count = 0
 119         result = 0
 120         shift = 0
 121
 122         for byte in encoded:
 123             count += 1
 124             result |= (byte & 0x7f) << shift
 125
 126             if not byte & 0x80:
 127                 return DecodedArg(self, zigzag_decode(result), encoded[:count])
 128
 129             shift += 7
 130             if shift >= 64:
 131                 break
 132
 133         return DecodedArg(self, None, encoded[:count], DecodedArg.DECODE_ERROR,
 134                           'Unterminated variable-length integer')
 135
 136     def _decode_unsigned_integer(self, encoded: bytes) -> 'DecodedArg':
 137         arg = self._decode_signed_integer(encoded)
 138
 139         # Since ZigZag encoding is used, unsigned integers must be masked off to
 140         # their original bit length.
 141         if arg.value is not None:
 142             arg.value &= (1 << self.size_bits()) - 1
 143
 144         return arg
 145
 146     def _decode_float(self, encoded: bytes) -> 'DecodedArg':
 147         if len(encoded) < 4:
 148             return DecodedArg.missing(self)
 149
 150         return DecodedArg(self,
 151                           self._PACKED_FLOAT.unpack_from(encoded)[0],
 152                           encoded[:4])
 153
 154     def _decode_string(self, encoded: bytes) -> 'DecodedArg':
 155         """Reads a unicode string from the encoded data."""
 156         if not encoded:
 157             return DecodedArg.missing(self)
 158
 159         size_and_status = encoded[0]
 160         status = DecodedArg.OK
 161
 162         if size_and_status & 0x80:
 163             status |= DecodedArg.TRUNCATED
 164             size_and_status &= 0x7f
 165
 166         raw_data = encoded[0:size_and_status + 1]
 167         data = raw_data[1:]
 168
 169         if len(data) < size_and_status:
 170             status |= DecodedArg.DECODE_ERROR
 171
 172         try:
 173             decoded = data.decode()
 174         except UnicodeDecodeError as err:
 175             return DecodedArg(self,
 176                               repr(bytes(data)).lstrip('b'), raw_data,
 177                               status | DecodedArg.DECODE_ERROR, err)
 178
 179         return DecodedArg(self, decoded, raw_data, status)
 180
 181     def _decode_char(self, encoded: bytes) -> 'DecodedArg':
 182         """Reads an integer from the data, then converts it to a string."""
 183         arg = self._decode_signed_integer(encoded)
 184
 185         if arg.ok():
 186             try:
 187                 arg.value = chr(arg.value)
 188             except (OverflowError, ValueError) as err:
 189                 arg.error = err
 190                 arg.status |= DecodedArg.DECODE_ERROR
 191
 192         return arg
 193
 194     def size_bits(self) -> int:
 195         """Size of the argument in bits; 0 for strings."""
 196         if self.type == 's':
 197             return 0
 198
 199         # TODO(hepler): 64-bit targets likely have 64-bit l, j, z, and t.
 200         return 64 if self.length in ['ll', 'j'] else 32
 201
 202     def __str__(self) -> str:
 203         return self.specifier
 204
 205
 206 class DecodedArg:
 207     """Represents a decoded argument that is ready to be formatted."""
 208
 209     # Status flags for a decoded argument. These values should match the
 210     # DecodingStatus enum in pw_tokenizer/internal/decode.h.
 211     OK = 0  # decoding was successful
 212     MISSING = 1  # the argument was not present in the data
 213     TRUNCATED = 2  # the argument was truncated during encoding
 214     DECODE_ERROR = 4  # an error occurred while decoding the argument
 215     SKIPPED = 8  # argument was skipped due to a previous error
 216
 217     @classmethod
 218     def missing(cls, specifier: FormatSpec):
 219         return cls(specifier, None, b'', cls.MISSING)
 220
 221     def __init__(self,
 222                  specifier: FormatSpec,
 223                  value,
 224                  raw_data: bytes,
 225                  status: int = OK,
 226                  error=None):
 227         self.specifier = specifier  # FormatSpec (e.g. to represent "%0.2f")
 228         self.value = value  # the decoded value, or None if decoding failed
 229         self.raw_data = bytes(
 230             raw_data)  # the exact bytes used to decode this arg
 231         self._status = status
 232         self.error = error
 233
 234     def ok(self) -> bool:
 235         """The argument was decoded without errors."""
 236         return self.status == self.OK or self.status == self.TRUNCATED
 237
 238     @property
 239     def status(self) -> int:
 240         return self._status
 241
 242     @status.setter
 243     def status(self, status: int):
 244         # The %% specifier is always OK and should always be printed normally.
 245         self._status = status if self.specifier.type != '%' else self.OK
 246
 247     def format(self) -> str:
 248         """Returns formatted version of this argument, with error handling."""
 249         if self.status == self.TRUNCATED:
 250             return self.specifier.compatible % (self.value + '[...]')
 251
 252         if self.ok():
 253             try:
 254                 return self.specifier.compatible % self.value
 255             except (OverflowError, TypeError, ValueError) as err:
 256                 self.status |= self.DECODE_ERROR
 257                 self.error = err
 258
 259         if self.status & self.SKIPPED:
 260             message = '{} SKIPPED'.format(self.specifier)
 261         elif self.status == self.MISSING:
 262             message = '{} MISSING'.format(self.specifier)
 263         elif self.status & self.DECODE_ERROR:
 264             message = '{} ERROR'.format(self.specifier)
 265         else:
 266             raise AssertionError('Unhandled DecodedArg status {:x}!'.format(
 267                 self.status))
 268
 269         if self.value is None or not str(self.value):
 270             return '<[{}]>'.format(message)
 271
 272         return '<[{} ({})]>'.format(message, self.value)
 273
 274     def __str__(self) -> str:
 275         return self.format()
 276
 277     def __repr__(self) -> str:
 278         return 'DecodedArg({!r})'.format(self)
 279
 280
 281 def parse_format_specifiers(format_string: str) -> Iterable[FormatSpec]:
 282     for spec in FormatSpec.FORMAT_SPEC.finditer(format_string):
 283         yield FormatSpec(spec)
 284
 285
 286 class FormattedString(NamedTuple):
 287     value: str
 288     args: Sequence[DecodedArg]
 289     remaining: bytes
 290
 291
 292 class FormatString:
 293     """Represents a printf-style format string."""
 294     def __init__(self, format_string: str):
 295         """Parses format specifiers in the format string."""
 296         self.format_string = format_string
 297         self.specifiers = tuple(parse_format_specifiers(self.format_string))
 298
 299         # List of non-specifier string pieces with room for formatted arguments.
 300         self._segments = self._parse_string_segments()
 301
 302     def _parse_string_segments(self) -> List:
 303         """Splits the format string by format specifiers."""
 304         if not self.specifiers:
 305             return [self.format_string]
 306
 307         spec_spans = [spec.match.span() for spec in self.specifiers]
 308
 309         # Start with the part of the format string up to the first specifier.
 310         string_pieces = [self.format_string[:spec_spans[0][0]]]
 311
 312         for ((_, end1), (start2, _)) in zip(spec_spans[:-1], spec_spans[1:]):
 313             string_pieces.append(self.format_string[end1:start2])
 314
 315         # Append the format string segment after the last format specifier.
 316         string_pieces.append(self.format_string[spec_spans[-1][1]:])
 317
 318         # Make a list with spots for the replacements between the string pieces.
 319         segments: List = [None] * (len(string_pieces) + len(self.specifiers))
 320         segments[::2] = string_pieces
 321
 322         return segments
 323
 324     def decode(self, encoded: bytes) -> Tuple[Sequence[DecodedArg], bytes]:
 325         """Decodes arguments according to the format string.
 326
 327         Args:
 328           encoded: bytes; the encoded arguments
 329
 330         Returns:
 331           tuple with the decoded arguments and any unparsed data
 332         """
 333         decoded_args = []
 334
 335         fatal_error = False
 336         index = 0
 337
 338         for spec in self.specifiers:
 339             arg = spec.decode(encoded[index:])
 340
 341             if fatal_error:
 342                 # After an error is encountered, continue to attempt to parse
 343                 # arguments, but mark them all as SKIPPED. If an error occurs,
 344                 # it's impossible to know if subsequent arguments are valid.
 345                 arg.status |= DecodedArg.SKIPPED
 346             elif not arg.ok():
 347                 fatal_error = True
 348
 349             decoded_args.append(arg)
 350             index += len(arg.raw_data)
 351
 352         return tuple(decoded_args), encoded[index:]
 353
 354     def format(self,
 355                encoded_args: bytes,
 356                show_errors: bool = False) -> FormattedString:
 357         """Decodes arguments and formats the string with them.
 358
 359         Args:
 360           encoded_args: the arguments to decode and format the string with
 361           show_errors: if True, an error message is used in place of the %
 362               conversion specifier when an argument fails to decode
 363
 364         Returns:
 365           tuple with the formatted string, decoded arguments, and remaining data
 366         """
 367         # Insert formatted arguments in place of each format specifier.
 368         args, remaining = self.decode(encoded_args)
 369
 370         if show_errors:
 371             self._segments[1::2] = (arg.format() for arg in args)
 372         else:
 373             self._segments[1::2] = (arg.format()
 374                                     if arg.ok() else arg.specifier.specifier
 375                                     for arg in args)
 376
 377         return FormattedString(''.join(self._segments), args, remaining)
 378
 379
 380 def decode(format_string: str,
 381            encoded_arguments: bytes,
 382            show_errors: bool = False) -> str:
 383     """Decodes arguments and formats them with the provided format string.
 384
 385     Args:
 386       format_string: the printf-style format string
 387       encoded_arguments: encoded arguments with which to format
 388           format_string; must exclude the 4-byte string token
 389       show_errors: if True, an error message is used in place of the %
 390           conversion specifier when an argument fails to decode
 391
 392     Returns:
 393       the printf-style formatted string
 394     """
 395     return FormatString(format_string).format(encoded_arguments,
 396                                               show_errors).value