[tor-commits] [stem/master] Add a stem.Endpoint class

[atagar at torproject.org]

commit dac2c5092c8f95a14bacc657297fc80a0207fd93
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Apr 20 11:58:22 2018 -0700

    Add a stem.Endpoint class
    
    Not yet used anywhere, but adding a little endpoint class so we can better
    differentiate ORPort vs DirPort endpoint arguments.
---
 stem/__init__.py        | 42 ++++++++++++++++++++++++++++++++++++++++++
 stem/util/connection.py |  2 ++
 test/settings.cfg       |  3 ++-
 test/unit/endpoint.py   | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/stem/__init__.py b/stem/__init__.py
index 853d704c..c8c43f65 100644
--- a/stem/__init__.py
+++ b/stem/__init__.py
@@ -618,6 +618,45 @@ class SocketClosed(SocketError):
   'Control socket was closed before completing the message.'
 
 
+class Endpoint(object):
+  """
+  Tor endpint that can be connected to.
+
+  .. versionadded:: 1.7.0
+
+  :var str address: ip address of the endpoint
+  :var int port: port of the endpoint
+  """
+
+  def __init__(self, address, port):
+    if not stem.util.connection.is_valid_ipv4_address(address) and not stem.util.connection.is_valid_ipv6_address(address):
+      raise ValueError("'%s' isn't a valid IPv4 or IPv6 address" % address)
+    elif not stem.util.connection.is_valid_port(port):
+      raise ValueError("'%s' isn't a valid port" % port)
+
+    self.address = address
+    self.port = int(port)
+
+
+class ORPort(Endpoint):
+  """
+  Tor relay's ORPort. The endpoint on which Tor accepts relay traffic.
+
+  :var list link_protocols: link protocol version we're willing to establish
+  """
+
+  def __init__(self, address, port, link_protocols = None):
+    super(ORPort, self).__init__(address, port)
+    self.link_protocols = link_protocols
+
+
+class DirPort(Endpoint):
+  """
+  Tor relay's DirPort. The endpoint on which Tor provides http access for
+  downloading descriptors.
+  """
+
+
 Runlevel = stem.util.enum.UppercaseEnum(
   'DEBUG',
   'INFO',
@@ -863,3 +902,6 @@ HSAuth = stem.util.enum.UppercaseEnum(
   'STEALTH_AUTH',
   'UNKNOWN',
 )
+
+
+import stem.util.connection  # importing afterward to avoid circular dependency
diff --git a/stem/util/connection.py b/stem/util/connection.py
index 3395a3a5..74fbdde2 100644
--- a/stem/util/connection.py
+++ b/stem/util/connection.py
@@ -439,6 +439,8 @@ def is_valid_ipv6_address(address, allow_brackets = False):
 
   if isinstance(address, bytes):
     address = str_tools._to_unicode(address)
+  elif not isinstance(address, str_type):
+    return False
 
   if allow_brackets:
     if address.startswith('[') and address.endswith(']'):
diff --git a/test/settings.cfg b/test/settings.cfg
index aaab0790..f1cbd381 100644
--- a/test/settings.cfg
+++ b/test/settings.cfg
@@ -132,7 +132,7 @@ pycodestyle.ignore E127
 pycodestyle.ignore E131
 pycodestyle.ignore E722
 
-pycodestyle.ignore stem/__init__.py => E402: import stem.util.enum
+pycodestyle.ignore stem/__init__.py => E402: import stem.util.connection
 pycodestyle.ignore stem/descriptor/__init__.py => E402: import stem.descriptor.server_descriptor
 pycodestyle.ignore stem/descriptor/__init__.py => E402: import stem.descriptor.extrainfo_descriptor
 pycodestyle.ignore stem/descriptor/__init__.py => E402: import stem.descriptor.networkstatus
@@ -216,6 +216,7 @@ test.unit_tests
 |test.unit.descriptor.certificate.TestEd25519Certificate
 |test.unit.exit_policy.rule.TestExitPolicyRule
 |test.unit.exit_policy.policy.TestExitPolicy
+|test.unit.endpoint.TestEndpoint
 |test.unit.version.TestVersion
 |test.unit.manual.TestManual
 |test.unit.tutorial.TestTutorial
diff --git a/test/unit/endpoint.py b/test/unit/endpoint.py
new file mode 100644
index 00000000..ad710a22
--- /dev/null
+++ b/test/unit/endpoint.py
@@ -0,0 +1,32 @@
+"""
+Unit tests for the stem.Endpoint class.
+"""
+
+import unittest
+
+import stem
+
+
+class TestEndpoint(unittest.TestCase):
+  def test_constructor(self):
+    endpoint = stem.ORPort('12.34.56.78', 80)
+    self.assertEqual('12.34.56.78', endpoint.address)
+    self.assertEqual(80, endpoint.port)
+    self.assertEqual(None, endpoint.link_protocols)
+
+    endpoint = stem.ORPort('12.34.56.78', 80, [3])
+    self.assertEqual('12.34.56.78', endpoint.address)
+    self.assertEqual(80, endpoint.port)
+    self.assertEqual([3], endpoint.link_protocols)
+
+    endpoint = stem.DirPort('12.34.56.78', 80)
+    self.assertEqual('12.34.56.78', endpoint.address)
+    self.assertEqual(80, endpoint.port)
+
+  def test_validation(self):
+    self.assertRaises(ValueError, stem.DirPort, '12.34.56.78', 'hello')
+    self.assertRaises(ValueError, stem.DirPort, '12.34.56.78', -5)
+    self.assertRaises(ValueError, stem.DirPort, '12.34.56.78', None)
+    self.assertRaises(ValueError, stem.DirPort, 'hello', 80)
+    self.assertRaises(ValueError, stem.DirPort, -5, 80)
+    self.assertRaises(ValueError, stem.DirPort, None, 80)