[tor-commits] [stem/master] Adding a is_private_address() function to stem.util.connection

[atagar at torproject.org]

commit aee7de80db712ab1003eb6d8ac52e78d5f0f18cc
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Sep 21 17:46:02 2013 -0700

    Adding a is_private_address() function to stem.util.connection
    
    Helper function for determining if an IPv4 address belongs to a private range
    or not.
---
 docs/change_log.rst          |    1 +
 stem/util/connection.py      |   35 +++++++++++++++++++++++++++++++++++
 test/unit/util/connection.py |   20 ++++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index b7840b4..13b256d 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -65,6 +65,7 @@ The following are only available within stem's `git repository
   * Added :func:`stem.util.system.get_user`
   * Added :func:`stem.util.system.get_start_time`
   * Added :func:`stem.util.system.get_bsd_jail_path`
+  * Added :func:`stem.util.connection.is_private_address`
 
  * **Website**
 
diff --git a/stem/util/connection.py b/stem/util/connection.py
index 3bb7df0..0e4c190 100644
--- a/stem/util/connection.py
+++ b/stem/util/connection.py
@@ -12,6 +12,7 @@ but for now just moving the parts we need.
   is_valid_ipv4_address - checks if a string is a valid IPv4 address
   is_valid_ipv6_address - checks if a string is a valid IPv6 address
   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
   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
@@ -119,6 +120,40 @@ def is_valid_port(entry, allow_zero = False):
   return entry > 0 and entry < 65536
 
 
+def is_private_address(address):
+  """
+  Checks if the IPv4 address is in a range belonging to the local network or
+  loopback. These include:
+
+    * Private ranges: 10.*, 172.16.* - 172.31.*, 192.168.*
+    * Loopback: 127.*
+
+  :param str address: string to be checked
+
+  :returns: **True** if input is in a private range, **False** otherwise
+
+  :raises: **ValueError** if the address isn't a valid IPv4 address
+  """
+
+  if not is_valid_ipv4_address(address):
+    raise ValueError("'%s' isn't a valid IPv4 address" % address)
+
+  # checks for any of the simple wildcard ranges
+
+  if address.startswith("10.") or address.startswith("192.168.") or address.startswith("127."):
+    return True
+
+  # checks for the 172.16.* - 172.31.* range
+
+  if address.startswith("172."):
+    second_octet = int(address.split('.')[1])
+
+    if second_octet >= 16 and second_octet <= 31: 
+      return True
+
+  return False
+
+
 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 99ee050..5d48f28 100644
--- a/test/unit/util/connection.py
+++ b/test/unit/util/connection.py
@@ -79,6 +79,26 @@ class TestConnection(unittest.TestCase):
     self.assertTrue(stem.util.connection.is_valid_port(0, allow_zero = True))
     self.assertTrue(stem.util.connection.is_valid_port("0", allow_zero = True))
 
+  def test_is_private_address(self):
+    """
+    Checks the is_private_address function.
+    """
+
+    self.assertTrue(stem.util.connection.is_private_address('127.0.0.1'))
+    self.assertTrue(stem.util.connection.is_private_address('10.0.0.0'))
+    self.assertTrue(stem.util.connection.is_private_address('172.16.0.0'))
+    self.assertTrue(stem.util.connection.is_private_address('172.31.0.0'))
+    self.assertTrue(stem.util.connection.is_private_address('192.168.0.50'))
+
+    self.assertFalse(stem.util.connection.is_private_address('74.125.28.103'))
+    self.assertFalse(stem.util.connection.is_private_address('172.15.0.0'))
+    self.assertFalse(stem.util.connection.is_private_address('172.32.0.0'))
+
+    self.assertRaises(ValueError, stem.util.connection.is_private_address, '')
+    self.assertRaises(ValueError, stem.util.connection.is_private_address, 'blarg')
+    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_expand_ipv6_address(self):
     """
     Checks the expand_ipv6_address function.