From fec9b6ee9b467187ed04133ed585488a0dada8cf Mon Sep 17 00:00:00 2001
From: ReenigneArcher <42013603+ReenigneArcher@users.noreply.github.com>
Date: Sat, 27 May 2023 11:21:16 -0400
Subject: [PATCH] fix(helpers): properly check public status of IPv6 addresses
 (#266)

---
 pyra/helpers.py            | 16 +++++++++++-----
 tests/unit/test_helpers.py | 36 +++++++++++++++++++++++++-----------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/pyra/helpers.py b/pyra/helpers.py
index d64116200..086494625 100644
--- a/pyra/helpers.py
+++ b/pyra/helpers.py
@@ -9,6 +9,7 @@
 
 # standard imports
 import datetime
+import ipaddress
 import logging
 import re
 import os
@@ -165,12 +166,16 @@ def is_public_ip(host: str) -> bool:
     False
     """
     ip = is_valid_ip(address=get_ip(host=host))
-    if ip and ip.iptype() == 'PUBLIC':
-        return True
-    return False
+
+    # use built in ipaddress module to check if address is private since IPy does not work IPv6 addresses
+    if ip:
+        ip_obj = ipaddress.ip_address(address=ip)
+        return not ip_obj.is_private
+    else:
+        return False
 
 
-def get_ip(host: str) -> str:
+def get_ip(host: str) -> Optional[str]:
     """
     Get IP address from host name.
 
@@ -184,7 +189,7 @@ def get_ip(host: str) -> str:
     Returns
     -------
     str
-        IP address of host name.
+        IP address of host name if it is a valid ip address, otherwise ``None``.
 
     Examples
     --------
@@ -201,6 +206,7 @@ def get_ip(host: str) -> str:
             ip_address = socket.getaddrinfo(host=host, port=None)[0][4][0]
         except Exception:
             log.error(f"IP Checker :: Bad IP or hostname provided: {host}.")
+            return None
         else:
             log.debug(f"IP Checker :: Resolved {host} to {ip_address}.")
             return ip_address
diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py
index dd404f203..86d2b60bc 100644
--- a/tests/unit/test_helpers.py
+++ b/tests/unit/test_helpers.py
@@ -27,17 +27,31 @@ def test_get_logger():
 
 
 def test_is_public_ip():
-    """Test if ip address is pubic, then tests if ip address is private"""
-    # test public ip
-    address = 'www.google.com'
-    ip = helpers.get_ip(host=address)
-    public_ip = helpers.is_public_ip(host=ip)
-    assert public_ip
-
-    # test private ip
-    ip = '192.168.1.1'
-    public_ip = helpers.is_public_ip(host=ip)
-    assert not public_ip
+    """Test if ip address is pubic or private"""
+    # test public ips
+    public_addresses = [
+        'www.google.com',
+        '8.8.8.8',
+        '8.8.4.4',
+        '2001:4860:4860::8888',
+        '2001:4860:4860::8844',
+    ]
+
+    private_addresses = [
+        '192.168.1.1'
+    ]
+
+    address_types = [public_addresses, private_addresses]
+
+    for address_type in address_types:
+        for address in address_type:
+            ip = helpers.get_ip(host=address)
+            public_ip = helpers.is_public_ip(host=ip)
+
+            if address_type == public_addresses:
+                assert public_ip
+            else:
+                assert not public_ip
 
 
 def test_get_ip():