Use a flag instead of a separate read_unused function

Author: andrivet

src/asn1.py

@@ -97,6 +97,10 @@ class Classes(Asn1Enum):
     Context = 0x80
     Private = 0xc0
 
+class ReadFlags(IntEnum):
+    OnlyValue = 0x00
+    WithUnused = 0x01
+
 
 Tag = collections.namedtuple('Tag', 'nr typ cls')
 """
@@ -620,12 +624,12 @@ class Decoder(object):
 
     def __init__(self):  # type: () -> None
         """Constructor."""
-        self._stream = None                 # type: Union[io.RawIOBase, None]  # Input stream
-        self._byte = bytes()                # type: bytes # Cached byte (to be able to implement eof)
-        self._position = 0                  # type: int # Due to caching, tell does not give the right position
-        self._tag = None                    # type: Union[Tag, None] # Cached Tag (to be able to implement peek)
-        self._levels = 0                    # type: int # Number of recursive calls
-        self._enters = 0                    # type int # Number of enter calls without leave
+        self._stream = None                     # type: Union[io.RawIOBase, None]  # Input stream
+        self._byte = bytes()                    # type: bytes # Cached byte (to be able to implement eof)
+        self._position = 0                      # type: int # Due to caching, tell does not give the right position
+        self._tag = None                        # type: Union[Tag, None] # Cached Tag (to be able to implement peek)
+        self._levels = 0                        # type: int # Number of recursive calls
+        self._enters = 0                        # type int # Number of enter calls without leave
 
     def start(self, stream):  # type: (Union[io.RawIOBase, bytes]) -> None
         """
@@ -689,26 +693,7 @@ def peek(self):  # type: () -> Union[Tag, None]
         self._tag = self._decode_tag()
         return self._tag
 
-    def read(self):  # type: () -> Tuple[Union[Tag, None], Any]
-        """
-        This method decodes one ASN.1 tag from the input and returns it as a
-        ``(tag, value, unused)`` tuple. ``tag`` is a 3-tuple ``(nr, typ, cls)``,
-        while ``value`` is a Python object representing the ASN.1 value.
-        The offset in the input is increased so that the next `Decoder.read()`
-        or `Decoder.read_with_unused()` call will return the next tag. In case
-        no more data is available from the input, this method returns ``None``
-        to signal end-of-file.
-
-        Returns:
-            `Tag`, value: The current ASN.1 tag and its value.
-
-        Raises:
-            `Error`
-        """
-        tag, value, _ = self.read_unused()
-        return tag, value
-
-    def read_unused(self):  # type: () -> Tuple[Union[Tag, None], Any, int]
+    def read(self, flags=ReadFlags.OnlyValue):  # type: (ReadFlags) -> Tuple[Union[Tag, None], Any]
         """
         This method decodes one ASN.1 tag from the input and returns it as a
         ``(tag, value, unused)`` tuple. ``tag`` is a 3-tuple ``(nr, typ, cls)``,
@@ -719,9 +704,12 @@ def read_unused(self):  # type: () -> Tuple[Union[Tag, None], Any, int]
         no more data is available from the input, this method returns ``None``
         to signal end-of-file.
 
+        Args:
+            flags (DecoderFlags): Return only the value or the value and the number of unused bits as a tuple.
+
         Returns:
-            `Tag`, value, unused: The current ASN.1 tag, its value and the
-            number of unused bits.
+            `Tag`, value: The current ASN.1 tag, its value and the
+            number of unused bits if requested in `start`.
 
         Raises:
             `Error`
@@ -731,10 +719,10 @@ def read_unused(self):  # type: () -> Tuple[Union[Tag, None], Any, int]
         tag = self.peek()
         self._tag = None
         if tag is None:
-            return None, None, 0
+            return None, None
         length = self._decode_length(tag.typ)
-        value, unused = self._decode_value(tag.cls, tag.typ, tag.nr, length)
-        return tag, value, unused
+        value = self._decode_value(tag.cls, tag.typ, tag.nr, length, flags)
+        return tag, value
 
     def eof(self):  # type: () -> bool
         """
@@ -889,7 +877,7 @@ def _decode_length(self, typ):  # type: (int) -> int
             length = byte
         return length
 
-    def _decode_value(self, cls, typ, nr, length):  # type: (int, int, int, int) -> Tuple[Any, int]
+    def _decode_value(self, cls, typ, nr, length, flags):  # type: (int, int, int, int, ReadFlags) -> Any
         """Read a value from the input. length is -1 in case of an indefinite form."""
         # A primitive type shall have a definite form
         if typ == Types.Primitive and length == INDEFINITE_FORM:
@@ -909,35 +897,38 @@ def _decode_value(self, cls, typ, nr, length):  # type: (int, int, int, int) ->
             raise Error('ASN1 decoding error: the Universal tag number {} shall have a constructed encoding'.format(nr))
 
         if cls != Classes.Universal:
-            return self._decode_bytes(typ, 0, length), 0
+            return self._decode_bytes(typ, 0, length)
 
         # Primitive encoding
         if nr == Numbers.Boolean:
-            return self._decode_boolean(length), 0
+            return self._decode_boolean(length)
         if nr in (Numbers.Integer, Numbers.Enumerated):
-            return self._decode_integer(length), 0
+            return self._decode_integer(length)
         if nr == Numbers.Null:
-            return self._decode_null(length), 0
+            return self._decode_null(length)
         if nr == Numbers.Real:
-            return self._decode_real(length), 0
+            return self._decode_real(length)
         if nr == Numbers.ObjectIdentifier:
-            return self._decode_object_identifier(length), 0
+            return self._decode_object_identifier(length)
 
         # Primitive or Constructed encoding
         if nr == Numbers.OctetString:
-            return self._decode_octet_string(typ, length), 0
+            return self._decode_octet_string(typ, length)
         if nr in (Numbers.PrintableString, Numbers.IA5String,
                   Numbers.UTF8String, Numbers.UTCTime,
                   Numbers.GeneralizedTime):
-            return self._decode_printable_string(typ, nr, length), 0
+            return self._decode_printable_string(typ, nr, length)
         if nr == Numbers.BitString:
-            return self._decode_bitstring(typ, length)
+            value, unused = self._decode_bitstring(typ, length, flags)
+            if flags & ReadFlags.WithUnused:
+                return value, unused
+            return value
 
         # Constructed types
         if nr == Numbers.Sequence:
-            return self._decode_sequence(length), 0
+            return self._decode_sequence(length)
 
-        return self._decode_bytes(typ, 0, length), 0
+        return self._decode_bytes(typ, 0, length)
 
     @staticmethod
     def _check_length(actual_length, expected_length):  # type: (int, int) -> None
@@ -1101,9 +1092,9 @@ def _is_eoc(self):  # type: () -> bool
         self._tag = None
         return True
 
-    def _decode_bitstring(self, typ, length):  # type: (int, int) -> Tuple[bytes, int]
+    def _decode_bitstring(self, typ, length, flags):  # type: (int, int, ReadFlags) -> Tuple[bytes, int]
         """Decode a bitstring."""
-        data, unused = self._decode_bytes_constructed(Numbers.BitString, length) \
+        data, unused = self._decode_bytes_constructed(Numbers.BitString, length, flags) \
             if typ == Types.Constructed else self._decode_bitstring_primitive(length)
         if self._levels == 0:
             data = shift_bits_right(data, unused)
@@ -1123,46 +1114,49 @@ def _decode_bitstring_primitive(self, length):  # type: (int) -> Tuple[bytes, in
 
         return content[1:], num_unused_bits
 
-    def _decode_bytes_constructed(self, nr, length):   # type: (int, int) -> Tuple[bytes, int]
+    def _decode_bytes_constructed(self, nr, length, flags):   # type: (int, int, ReadFlags) -> Tuple[bytes, int]
         if length == INDEFINITE_FORM:
-            return self._decode_bytes_indefinite(nr)
+            return self._decode_bytes_indefinite(nr, flags)
         else:
-            return self._decode_bytes_definite(nr, length)
+            return self._decode_bytes_definite(nr, length, flags)
 
-    def _decode_bytes_indefinite(self, nr):  # type: (int) -> Tuple[bytes, int]
+    def _decode_bytes_indefinite(self, nr, flags):  # type: (int, ReadFlags) -> Tuple[bytes, int]
         self._levels += 1
         value = b''
         unused = 0
         while not self._is_eoc():
             if unused > 0:
                 raise Error('ASN1 decoding error: unused bits shall be 0 unless it is the last segment')
-            segment_value, segment_unused = self._decode_bytes_segment(nr)
+            segment_value, segment_unused = self._decode_bytes_segment(nr, flags)
             value += segment_value
             unused = segment_unused
         self._levels -= 1
         return value, unused
 
-    def _decode_bytes_segment(self, nr):  # type: (int) -> Tuple[bytes, int]
+    def _decode_bytes_segment(self, nr, flags):  # type: (int, ReadFlags) -> Tuple[bytes, int]
         tag = self.peek()
         if tag is None:
             raise Error("ASN1 decoding error: premature end of input.")
         if nr != 0 and (tag.nr != nr or tag.cls != Classes.Universal):
             raise Error(
                 'ASN1 decoding error: invalid tag (cls={}, nr={}) in a constructed type.'.format(tag.cls, tag.nr))
-        tag, value, unused = self.read_unused()
+        tag, value = self.read(flags)
+        unused = 0
+        if flags & ReadFlags.WithUnused:
+            value, unused = value
         if not isinstance(value, bytes):
             raise Error('ASN1 decoding error: bytes were expected instead of the type {}'.format(type(value)))
         return value, unused
 
-    def _decode_bytes_definite(self, nr, length):  # type: (int, int) -> Tuple[bytes, int]
+    def _decode_bytes_definite(self, nr, length, flags):  # type: (int, int, ReadFlags) -> Tuple[bytes, int]
         self._levels += 1
         value = b''
         unused = 0
         start_position = self._get_current_position()
         while self._get_current_position() - start_position < length:
             if unused > 0:
                 raise Error('ASN1 decoding error: unused bits shall be 0 unless it is the last segment')
-            segment_value, segment_unused = self._decode_bytes_segment(nr)
+            segment_value, segment_unused = self._decode_bytes_segment(nr, flags)
             value += segment_value
             unused = segment_unused
 
@@ -1176,15 +1170,15 @@ def _decode_octet_string(self, typ, length):  # type: (int, int) -> bytes
         """Decode an octet string."""
         if typ == Types.Primitive:
             return self._read_bytes(length)
-        value, unused = self._decode_bytes_constructed(Numbers.OctetString, length)
+        value, unused = self._decode_bytes_constructed(Numbers.OctetString, length, ReadFlags.OnlyValue)
         assert unused == 0
         return value
 
     def _decode_bytes(self, typ, nr, length):  # type: (int, int, int) -> bytes
         """Decode a string."""
         if typ == Types.Primitive:
             return self._read_bytes(length)
-        value, unused = self._decode_bytes_constructed(nr, length)
+        value, unused = self._decode_bytes_constructed(nr, length, ReadFlags.OnlyValue)
         assert unused == 0
         return value
 

tests/test_asn1.py

@@ -639,7 +639,7 @@ def test_bitstring_primitive(self):
         dec.start(buf)
         tag = dec.peek()
         assert tag == (asn1.Numbers.BitString, asn1.Types.Primitive, asn1.Classes.Universal)
-        tag, val, unused = dec.read_unused()
+        tag, (val, unused) = dec.read(asn1.ReadFlags.WithUnused)
         assert val == b'\x12\x34\x56'
         assert unused == 0
 
@@ -649,7 +649,7 @@ def test_bitstring_constructed(self):
         dec.start(buf)
         tag = dec.peek()
         assert tag == (asn1.Numbers.BitString, asn1.Types.Constructed, asn1.Classes.Universal)
-        tag, val, unused = dec.read_unused()
+        tag, (val, unused) = dec.read(asn1.ReadFlags.WithUnused)
         assert val == b'\x00\xB0\xB0'
         assert unused == 4
 
@@ -659,7 +659,7 @@ def test_bitstring_unused_bits(self):
         dec.start(buf)
         tag = dec.peek()
         assert tag == (asn1.Numbers.BitString, asn1.Types.Primitive, asn1.Classes.Universal)
-        ttag, val, unused = dec.read_unused()
+        tag, (val, unused) = dec.read(asn1.ReadFlags.WithUnused)
         assert val == b'\x01\x23\x45'
         assert unused == 4
 

tests/test_suite.py

@@ -385,15 +385,15 @@ def test_case_36(self, decode_ber):
         Using of "unused bits" in internal BIT STRINGs with constructive form of encoding.
         """
         with pytest.raises(asn1.Error) as e:
-            decode_ber.read()
+            decode_ber.read(asn1.ReadFlags.WithUnused)
         assert "unused bits shall be 0" in str(e.value)
 
     @pytest.mark.parametrize('filename', ['tc37.ber'])
     def test_case_37(self, decode_ber):
         """
         Using of definite form of length block in case of constructive form of encoding.
         """
-        tag, value, unused = decode_ber.read_unused()
+        tag, (value, unused) = decode_ber.read(asn1.ReadFlags.WithUnused)
         assert tag == (asn1.Numbers.BitString, asn1.Types.Constructed, asn1.Classes.Universal)
         assert isinstance(value, bytes)
         assert value == b'\x00\x10\x10'
@@ -404,7 +404,7 @@ def test_case_38(self, decode_ber):
         """
         Using of indefinite form of length block in case of constructive form of encoding.
         """
-        tag, value, unused = decode_ber.read_unused()
+        tag, (value, unused) = decode_ber.read(asn1.ReadFlags.WithUnused)
         assert tag == (asn1.Numbers.BitString, asn1.Types.Constructed, asn1.Classes.Universal)
         assert isinstance(value, bytes)
         assert value == b'\x00\xA3\xB5\xF2\x91\xCD'
@@ -415,7 +415,7 @@ def test_case_39(self, decode_ber):
         """
         Using of constructive form of encoding for empty BIT STRING.
         """
-        tag, value, unused = decode_ber.read_unused()
+        tag, (value, unused) = decode_ber.read(asn1.ReadFlags.WithUnused)
         assert tag == (asn1.Numbers.BitString, asn1.Types.Constructed, asn1.Classes.Universal)
         assert isinstance(value, bytes)
         assert value == b''
@@ -429,7 +429,7 @@ def test_case_40(self, decode_ber):
         So this case raises an exception
         """
         with pytest.raises(asn1.Error) as e:
-            decode_ber.read_unused()
+            decode_ber.read(asn1.ReadFlags.WithUnused)
         assert "initial byte is missing" in str(e.value)
 
 class TestSuiteOctetString: