Use TLSPrefixedArray in tls.message.Certificate.from_bytes (#128)

* Use `TLSPrefixedArray` in `tls.message.Certificate.from_bytes` Fixes #109. * Fixing errors * Finally removing BytesIO
python-tls · Jan 30, 2017 · 6106a21 · 6106a21
1 parent ad06bce
commit 6106a21
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 26 deletions.
diff --git a/tls/_constructs.py b/tls/_constructs.py
@@ -251,16 +251,18 @@
 
 ASN1Cert = Struct(
     "ASN1Cert",
-    SizeWithin(UBInt24("length"),
-               min_size=1, max_size=2 ** 24 - 1),
-    Bytes("asn1_cert", lambda ctx: ctx.length),
+    PrefixedBytes("asn1_cert", SizeWithin(UBInt24("length"), min_size=1,
+                                          max_size=2 ** 24 - 1))
 )
 
+# https://tools.ietf.org/html/rfc5246#section-7.4.2
 Certificate = Struct(
     "Certificate",
-    SizeWithin(UBInt24("certificates_length"),
-               min_size=1, max_size=2 ** 24 - 1),
-    Bytes("certificates_bytes", lambda ctx: ctx.certificates_length),
+    TLSPrefixedArray("certificate_list",
+                     ASN1Cert,
+                     length_validator=partial(SizeWithin, min_size=1,
+                                              max_size=2 ** 24 - 1),
+                     length_field_size=UBInt24),
 )
 
 Handshake = Struct(

diff --git a/tls/message.py b/tls/message.py
@@ -8,8 +8,6 @@
 
 from construct import Container, ListContainer
 
-from six import BytesIO
-
 from tls import _constructs
 
 from tls._common import enums
@@ -171,12 +169,8 @@ class Certificate(object):
 
     def as_bytes(self):
         return _constructs.Certificate.build(Container(
-            certificates_length=sum([3 + len(asn1cert.asn1_cert)
-                                     for asn1cert in self.certificate_list]),
-            certificates_bytes=b''.join(
-                [asn1cert.as_bytes() for asn1cert in self.certificate_list]
-            )
-
+            certificate_list=[Container(asn1_cert=cert.asn1_cert)
+                              for cert in self.certificate_list]
         ))
 
     @classmethod
@@ -188,19 +182,12 @@ def from_bytes(cls, bytes):
         :return: Certificate object.
         """
         construct = _constructs.Certificate.parse(bytes)
-        # XXX: Find a better way to parse an array of variable-length objects
-        certificates = []
-        certificates_io = BytesIO(construct.certificates_bytes)
-
-        while certificates_io.tell() < construct.certificates_length:
-            certificate_construct = _constructs.ASN1Cert.parse_stream(
-                certificates_io
-            )
-            certificates.append(
-                ASN1Cert(asn1_cert=certificate_construct.asn1_cert)
-            )
         return cls(
-            certificate_list=certificates
+            certificate_list=[
+                ASN1Cert(
+                    asn1_cert=asn1cert.asn1_cert
+                )
+                for asn1cert in construct.certificate_list],
         )