[tor-commits] [stem/master] Filter archives with multiple descriptor types

[atagar at torproject.org]

commit c54a66b92c63e9156b884c1dedcc3bd2da443f38
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Aug 9 16:40:07 2019 -0700

    Filter archives with multiple descriptor types
    
    One CollecTor archive type contains multiple descriptor types, so filtering out
    unrequested descriptors.
---
 stem/descriptor/collector.py           | 12 ++++++++----
 stem/descriptor/router_status_entry.py |  6 ++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/stem/descriptor/collector.py b/stem/descriptor/collector.py
index c8c4a1b7..3867fa3d 100644
--- a/stem/descriptor/collector.py
+++ b/stem/descriptor/collector.py
@@ -263,12 +263,16 @@ class File(object):
 
         return
 
-    # TODO: the following will not work if the tar contains multiple types or a type we do not support
-
     path = self.download(directory, True, timeout, retries)
 
-    for desc in parse_file(path, descriptor_type):
-      yield desc
+    # Archives can contain multiple descriptor types, so parsing everything and
+    # filtering to what we're after.
+
+    for desc in parse_file(path):
+      desc_annotation = type(desc).TYPE_ANNOTATION_NAME
+
+      if descriptor_type is None or (desc_annotation and descriptor_type.startswith(desc_annotation)):
+        yield desc
 
   def download(self, directory, decompress = True, timeout = None, retries = 3):
     """
diff --git a/stem/descriptor/router_status_entry.py b/stem/descriptor/router_status_entry.py
index 344c5697..31913417 100644
--- a/stem/descriptor/router_status_entry.py
+++ b/stem/descriptor/router_status_entry.py
@@ -515,6 +515,8 @@ class RouterStatusEntryV2(RouterStatusEntry):
   a default value, others are left as **None** if undefined
   """
 
+  TYPE_ANNOTATION_NAME = 'network-status-consensus-2'
+
   ATTRIBUTES = dict(RouterStatusEntry.ATTRIBUTES, **{
     'digest': (None, _parse_r_line),
   })
@@ -575,6 +577,8 @@ class RouterStatusEntryV3(RouterStatusEntry):
      Added the protocols attribute.
   """
 
+  TYPE_ANNOTATION_NAME = 'network-status-consensus-3'
+
   ATTRIBUTES = dict(RouterStatusEntry.ATTRIBUTES, **{
     'digest': (None, _parse_r_line),
     'or_addresses': ([], _parse_a_line),
@@ -652,6 +656,8 @@ class RouterStatusEntryMicroV3(RouterStatusEntry):
   a default value, others are left as **None** if undefined
   """
 
+  TYPE_ANNOTATION_NAME = 'network-status-microdesc-consensus-3'
+
   ATTRIBUTES = dict(RouterStatusEntry.ATTRIBUTES, **{
     'or_addresses': ([], _parse_a_line),
     'bandwidth': (None, _parse_w_line),