Give ip2long and long2ip ipv6 support (#271)

Improved ipv6 support in the `ip2long` and `long2ip` helper functions

CTFd/models.py

@@ -1,7 +1,7 @@
 import datetime
 import hashlib
 import json
-from socket import inet_aton, inet_ntoa
+from socket import inet_pton, inet_ntop, AF_INET, AF_INET6
 from struct import unpack, pack, error as struct_error
 
 from flask_sqlalchemy import SQLAlchemy
@@ -14,15 +14,24 @@ def sha512(string):
 
 
 def ip2long(ip):
-    return unpack('!i', inet_aton(ip))[0]
+    '''Converts a user's IP address into an integer/long'''
+    if '.' in ip:
+        # ipv4
+        return unpack('!i', inet_pton(AF_INET, ip))[0]
+    else:
+        # ipv6
+        hi, lo = unpack('!QQ', inet_pton(AF_INET6, ip))
+        return (hi << 64) | lo
 
 
 def long2ip(ip_int):
-    try:
+    '''Converts a saved IP address back into an integer/long'''
-        return inet_ntoa(pack('!i', ip_int))
+    if ip_int < 4294967296:
-    except struct_error:
+        # ipv4
-        # Backwards compatibility with old CTFd databases
+        return inet_ntop(AF_INET, pack('!i', ip_int))
-        return inet_ntoa(pack('!I', ip_int))
+    else:
+        # ipv6
+        return inet_ntop(AF_INET6, pack('!QQ', ip_int >> 64, ip_int & 0xffffffffffffffff))
 
 
 db = SQLAlchemy()

tests/test_utils.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from tests.helpers import *
+from CTFd.models import ip2long, long2ip
+import json
+
+
+def test_ip2long_ipv4():
+    """Does ip2long work properly for ipv4 addresses"""
+    assert ip2long('127.0.0.1') == 2130706433
+
+
+def test_long2ip_ipv4():
+    """Does long2ip work properly for ipv4 addresses"""
+    assert long2ip(2130706433) == '127.0.0.1'
+
+
+def test_ip2long_ipv6():
+    """Does ip2long work properly for ipv6 addresses"""
+    assert ip2long('2001:0db8:85a3:0000:0000:8a2e:0370:7334') == 42540766452641154071740215577757643572
+    assert ip2long('2001:658:22a:cafe:200::1') == 42540616829182469433547762482097946625
+
+
+def test_long2ip_ipv6():
+    """Does long2ip work properly for ipv6 addresses"""
+    assert long2ip(42540766452641154071740215577757643572) == '2001:db8:85a3::8a2e:370:7334'
+    assert long2ip(42540616829182469433547762482097946625) == '2001:658:22a:cafe:200::1'