[tor-commits] [stem/master] Being more anal about 'OK' status checking

[atagar at torproject.org]

commit 6015799baf9c0e6d3a618c948c32c27b2d7e1fc3
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon May 28 15:48:32 2012 -0700

    Being more anal about 'OK' status checking
    
    Refusing to parse messages that lack a 250 response on all lines, and checking
    for a 'OK' status on the last line (and only the last line).
---
 stem/response/getinfo.py      |   10 +++++++---
 stem/response/protocolinfo.py |   11 ++++++-----
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/stem/response/getinfo.py b/stem/response/getinfo.py
index a13a18f..7cde9f1 100644
--- a/stem/response/getinfo.py
+++ b/stem/response/getinfo.py
@@ -23,9 +23,13 @@ class GetInfoResponse(stem.socket.ControlMessage):
     
     self.values = {}
     
-    for line in self:
-      if line == "OK": break
-      elif not "=" in line:
+    lines = list(self)
+    
+    if not self.is_ok() or not lines.pop() == "OK":
+      raise stem.socket.ProtocolError("GETINFO response didn't have an OK status:\n%s" % self)
+    
+    for line in lines:
+      if not "=" in line:
         raise stem.socket.ProtocolError("GETINFO replies should only contain parameter=value mappings: %s" % line)
       
       key, value = line.split("=", 1)
diff --git a/stem/response/protocolinfo.py b/stem/response/protocolinfo.py
index 4547a1d..97ecfb7 100644
--- a/stem/response/protocolinfo.py
+++ b/stem/response/protocolinfo.py
@@ -37,16 +37,17 @@ class ProtocolInfoResponse(stem.socket.ControlMessage):
     self.cookie_path = None
     
     auth_methods, unknown_auth_methods = [], []
+    lines = list(self)
+    
+    if not self.is_ok() or not lines.pop() == "OK":
+      raise stem.socket.ProtocolError("GETINFO response didn't have an OK status:\n%s" % self)
     
     # sanity check that we're a PROTOCOLINFO response
-    if not list(self)[0].startswith("PROTOCOLINFO"):
+    if not lines[0].startswith("PROTOCOLINFO"):
       msg = "Message is not a PROTOCOLINFO response (%s)" % self
       raise stem.socket.ProtocolError(msg)
     
-    for line in self:
-      if line == "OK": break
-      elif line.is_empty(): continue # blank line
-      
+    for line in lines:
       line_type = line.pop()
       
       if line_type == "PROTOCOLINFO":