[tor-commits] [stem/master] Add stem.util.connection.address_to_int()

[atagar at torproject.org]

commit fd507a61eea24b99c260fc47552279f9deba6f18
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Jul 15 10:29:08 2016 -0700

    Add stem.util.connection.address_to_int()
    
    Couple times now I've needed to be able to serialize an IP address into an int.
    Might as well finally add a function for this. This can be used for realization
    if we had the other counterpart (int_to_address()), but not something I need
    right now.
---
 docs/change_log.rst          |  1 +
 stem/exit_policy.py          |  4 ++--
 stem/util/connection.py      | 20 ++++++++++++++++++++
 test/unit/util/connection.py | 11 +++++++++++
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index 12764e8..98dc18c 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -93,6 +93,7 @@ The following are only available within Stem's `git repository
   * Added **stem.util.system.SYSTEM_CALL_TIME** with the total time spent on system calls
   * Added an **is_ipv6** value to :class:`~stem.util.connection.Connection` instances
   * Added :func:`~stem.util.system.pids_by_user`
+  * Added :func:`~stem.util.connection.address_to_int`
   * Added :func:`~stem.util.__init__.datetime_to_unix`
 
  * **Interpreter**
diff --git a/stem/exit_policy.py b/stem/exit_policy.py
index 9379d9c..df6809f 100644
--- a/stem/exit_policy.py
+++ b/stem/exit_policy.py
@@ -765,7 +765,7 @@ class ExitPolicyRule(object):
       if address is None:
         fuzzy_match = True
       else:
-        comparison_addr_bin = int(stem.util.connection._get_address_binary(address), 2)
+        comparison_addr_bin = stem.util.connection.address_to_int(address)
         comparison_addr_bin &= self._get_mask_bin()
 
         if self._get_address_bin() != comparison_addr_bin:
@@ -907,7 +907,7 @@ class ExitPolicyRule(object):
   def _get_address_bin(self):
     # provides an integer representation of our address
 
-    return int(stem.util.connection._get_address_binary(self.address), 2) & self._get_mask_bin()
+    return stem.util.connection.address_to_int(self.address) & self._get_mask_bin()
 
   def _apply_addrspec(self, rule, addrspec, is_ipv6_only):
     # Parses the addrspec...
diff --git a/stem/util/connection.py b/stem/util/connection.py
index 49a410b..a375a17 100644
--- a/stem/util/connection.py
+++ b/stem/util/connection.py
@@ -17,6 +17,8 @@ Connection and networking based utility functions.
   is_valid_port - checks if something is a valid representation for a port
   is_private_address - checks if an IPv4 address belongs to a private range or not
 
+  address_to_int - provides an integer representation of an IP address
+
   expand_ipv6_address - provides an IPv6 address with its collapsed portions expanded
   get_mask_ipv4 - provides the mask representation for a given number of bits
   get_mask_ipv6 - provides the IPv6 mask representation for a given number of bits
@@ -523,6 +525,24 @@ def is_private_address(address):
   return False
 
 
+def address_to_int(address):
+  """
+  Provides an integer representation of a IPv4 or IPv6 address that can be used
+  for sorting.
+
+  .. versionadded:: 1.5.0
+
+  :param str address: IPv4 or IPv6 address
+
+  :returns: **int** representation of the address
+  """
+
+  # TODO: Could be neat to also use this for serialization if we also had an
+  # int_to_address() function.
+
+  return int(_get_address_binary(address), 2)
+
+
 def expand_ipv6_address(address):
   """
   Expands abbreviated IPv6 addresses to their full colon separated hex format.
diff --git a/test/unit/util/connection.py b/test/unit/util/connection.py
index e645c72..14d19a2 100644
--- a/test/unit/util/connection.py
+++ b/test/unit/util/connection.py
@@ -494,6 +494,17 @@ class TestConnection(unittest.TestCase):
     self.assertRaises(ValueError, stem.util.connection.is_private_address, '127.0.0')
     self.assertRaises(ValueError, stem.util.connection.is_private_address, 'fe80:0000:0000:0000:0202:b3ff:fe1e:8329')
 
+  def test_address_to_int(self):
+    """
+    Checks the address_to_int function.
+    """
+
+    self.assertEqual(1, stem.util.connection.address_to_int('0.0.0.1'))
+    self.assertEqual(2, stem.util.connection.address_to_int('0.0.0.2'))
+    self.assertEqual(256, stem.util.connection.address_to_int('0.0.1.0'))
+    self.assertEqual(2130706433, stem.util.connection.address_to_int('127.0.0.1'))
+    self.assertEqual(338288524927261089654163772891438416681L, stem.util.connection.address_to_int('fe80:0000:0000:0000:0202:b3ff:fe1e:8329'))
+
   def test_expand_ipv6_address(self):
     """
     Checks the expand_ipv6_address function.