[tor-commits] [stem/master] Supporting v2 documents in the descriptor reader

[atagar at torproject.org]

commit ec335d1b5bf4368982bc9736a271f675a6a7c60d
Author: Damian Johnson <atagar at torproject.org>
Date:   Thu Oct 11 19:54:05 2012 -0700

    Supporting v2 documents in the descriptor reader
    
    The v2 descriptor's integ data broke the descriptor reader tests. Turned that I
    had forgotten to add v2 documents to the types that it recognizes.
---
 stem/descriptor/__init__.py                    |    6 +++-
 stem/descriptor/networkstatus.py               |   25 ++++++++++++++++-------
 test/integ/descriptor/data/cached-consensus-v2 |    1 +
 test/integ/descriptor/networkstatus.py         |    1 +
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/stem/descriptor/__init__.py b/stem/descriptor/__init__.py
index f0ab82d..9d33ab7 100644
--- a/stem/descriptor/__init__.py
+++ b/stem/descriptor/__init__.py
@@ -25,7 +25,6 @@ __all__ = [
 
 import os
 import re
-import datetime
 import collections
 
 import stem.util.enum
@@ -86,7 +85,7 @@ def parse_file(path, descriptor_file):
   elif filename == "cached-consensus":
     file_parser = stem.descriptor.networkstatus.parse_file
   elif filename == "cached-microdesc-consensus":
-    file_parser = lambda f: stem.descriptor.networkstatus.parse_file(f, True, True)
+    file_parser = lambda f: stem.descriptor.networkstatus.parse_file(f, is_microdescriptor = True)
   else:
     # Metrics descriptor handling
     first_line, desc = descriptor_file.readline().strip(), None
@@ -128,6 +127,9 @@ def _parse_metrics_file(descriptor_type, major_version, minor_version, descripto
   elif descriptor_type in ("network-status-consensus-3", "network-status-vote-3") and major_version == 1:
     for desc in stem.descriptor.networkstatus.parse_file(descriptor_file):
       yield desc
+  elif descriptor_type == "network-status-2" and major_version == 1:
+    for desc in stem.descriptor.networkstatus.parse_file(descriptor_file, document_version = 2):
+      yield desc
   elif descriptor_type == "network-status-microdesc-consensus-3" and major_version == 1:
     for desc in stem.descriptor.networkstatus.parse_file(descriptor_file, is_microdescriptor = True):
       yield desc
diff --git a/stem/descriptor/networkstatus.py b/stem/descriptor/networkstatus.py
index a028572..21f7ee0 100644
--- a/stem/descriptor/networkstatus.py
+++ b/stem/descriptor/networkstatus.py
@@ -164,7 +164,7 @@ BANDWIDTH_WEIGHT_ENTRIES = (
   "Wmb", "Wmd", "Wme", "Wmg", "Wmm",
 )
 
-def parse_file(document_file, validate = True, is_microdescriptor = False):
+def parse_file(document_file, validate = True, is_microdescriptor = False, document_version = 3):
   """
   Parses a network status and iterates over the RouterStatusEntry in it. The
   document that these instances reference have an empty 'rotuers' attribute to
@@ -173,29 +173,38 @@ def parse_file(document_file, validate = True, is_microdescriptor = False):
   :param file document_file: file with network status document content
   :param bool validate: checks the validity of the document's contents if True, skips these checks otherwise
   :param bool is_microdescriptor: True if this is for a microdescriptor consensus, False otherwise
+  :param int document_version: network status document version
   
-  :returns: :class:`stem.descriptor.networkstatus.NetworkStatusDocumentV3` object
+  :returns: :class:`stem.descriptor.networkstatus.NetworkStatusDocument` object
   
   :raises:
-    * ValueError if the contents is malformed and validate is True
+    * ValueError if the document_version is unrecognized or the contents is malformed and validate is True
     * IOError if the file can't be read
   """
   
   # getting the document without the routers section
   
-  header = stem.descriptor._read_until_keywords((ROUTERS_START, FOOTER_START), document_file)
+  header = stem.descriptor._read_until_keywords((ROUTERS_START, FOOTER_START, V2_FOOTER_START), document_file)
   
   routers_start = document_file.tell()
-  stem.descriptor._read_until_keywords(FOOTER_START, document_file, skip = True)
+  stem.descriptor._read_until_keywords((FOOTER_START, V2_FOOTER_START), document_file, skip = True)
   routers_end = document_file.tell()
   
   footer = document_file.readlines()
   document_content = "".join(header + footer)
   
-  if not is_microdescriptor:
+  if document_version == 2:
+    document_type = NetworkStatusDocumentV2
     router_type = stem.descriptor.router_status_entry.RouterStatusEntryV3
+  elif document_version == 3:
+    document_type = NetworkStatusDocumentV3
+    
+    if not is_microdescriptor:
+      router_type = stem.descriptor.router_status_entry.RouterStatusEntryV3
+    else:
+      router_type = stem.descriptor.router_status_entry.RouterStatusEntryMicroV3
   else:
-    router_type = stem.descriptor.router_status_entry.RouterStatusEntryMicroV3
+    raise ValueError("Document version %i isn't recognized (only able to parse v2 or v3)" % document_version)
   
   desc_iterator = _get_entries(
     document_file,
@@ -204,7 +213,7 @@ def parse_file(document_file, validate = True, is_microdescriptor = False):
     entry_keyword = ROUTERS_START,
     start_position = routers_start,
     end_position = routers_end,
-    extra_args = (NetworkStatusDocumentV3(document_content, validate),),
+    extra_args = (document_type(document_content, validate),),
   )
   
   for desc in desc_iterator:
diff --git a/test/integ/descriptor/data/cached-consensus-v2 b/test/integ/descriptor/data/cached-consensus-v2
index 4f8cbce..ad64d40 100644
--- a/test/integ/descriptor/data/cached-consensus-v2
+++ b/test/integ/descriptor/data/cached-consensus-v2
@@ -1,3 +1,4 @@
+ at type network-status-2 1.0
 network-status-version 2
 dir-source 18.244.0.114 18.244.0.114 80
 fingerprint 719BE45DE224B607C53707D0E2143E2D423E74CF
diff --git a/test/integ/descriptor/networkstatus.py b/test/integ/descriptor/networkstatus.py
index e0ecd00..96ce12f 100644
--- a/test/integ/descriptor/networkstatus.py
+++ b/test/integ/descriptor/networkstatus.py
@@ -213,6 +213,7 @@ TpQQk3nNQF8z6UIvdlvP+DnJV4izWVkQEZgUZgIVM0E=
     consensus_path = test.integ.descriptor.get_resource("cached-consensus-v2")
     
     with open(consensus_path) as descriptor_file:
+      descriptor_file.readline() # strip header
       document = stem.descriptor.networkstatus.NetworkStatusDocumentV2(descriptor_file.read())
       
       self.assertEquals(2, document.version)