From ed8486e2a163d57ca8bb931422b98a16e5aea827 Mon Sep 17 00:00:00 2001
From: Prodesire <prodesire@yeah.net>
Date: Fri, 8 Dec 2017 22:54:47 +0800
Subject: [PATCH] add network.is_ipv6

---
 docs/network.rst      | 11 +++++++++++
 pydu/network.py       | 20 ++++++++++++++++++--
 tests/test_network.py | 12 +++++++++++-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/docs/network.rst b/docs/network.rst
index f187ed3..5fa2740 100644
--- a/docs/network.rst
+++ b/docs/network.rst
@@ -24,6 +24,17 @@ Network
     False
 
 
+.. py:function:: pydu.network.is_ipv6(ip)
+
+    Judge whether the given ``ip`` is IPV6 address.
+
+    >>> from pydu.network import is_ipv6
+    >>> is_ipv4('fe80::9e5b:b149:e187:1a18')
+    True
+    >>> is_ipv4('localhost.localdomain')
+    False
+
+
 .. py:function:: pydu.network.get_free_port()
 
     Get free port which could be bound.
diff --git a/pydu/network.py b/pydu/network.py
index e4479ca..8eb70bd 100644
--- a/pydu/network.py
+++ b/pydu/network.py
@@ -13,10 +13,26 @@ def dotted_netmask(mask):
     return socket.inet_ntoa(struct.pack('>I', bits))
 
 
-# https://github.com/kennethreitz/requests/blob/master/requests/utils.py
 def is_ipv4(ip):
+    """
+    Returns True if the IPv4 address ia valid, otherwise returns False.
+    """
+    return _is_ip(ip, af=socket.AF_INET)
+
+
+def is_ipv6(ip):
+    """
+    Returns True if the IPv6 address ia valid, otherwise returns False.
+    """
+    return _is_ip(ip, af=socket.AF_INET6)
+
+
+def _is_ip(ip, af):
+    '''
+    Returns True if the IP address ia valid, otherwise returns False.
+    '''
     try:
-        socket.inet_aton(ip)
+        socket.inet_pton(af, ip)
     except socket.error:
         return False
     return True
diff --git a/tests/test_network.py b/tests/test_network.py
index 4135ceb..335a714 100644
--- a/tests/test_network.py
+++ b/tests/test_network.py
@@ -1,5 +1,5 @@
 import pytest
-from pydu.network import dotted_netmask, is_ipv4, get_free_port
+from pydu.network import dotted_netmask, is_ipv4, is_ipv6, get_free_port
 
 
 @pytest.mark.parametrize(
@@ -22,6 +22,16 @@ class TestIsIPv4Address:
         assert not is_ipv4(value)
 
 
+class TestIsIPv6Address:
+
+    def test_valid(self):
+        assert is_ipv6('fe80::9e5b:b149:e187:1a18')
+
+    @pytest.mark.parametrize('value', ('fe80::9e5b:b149:e187::', 'localhost.localdomain'))
+    def test_invalid(self, value):
+        assert not is_ipv6(value)
+
+
 def test_get_free_port():
     port = get_free_port()
     assert isinstance(port, int)