diff --git a/CTFd/models.py b/CTFd/models.py index 4befa17c..daab74f2 100644 --- a/CTFd/models.py +++ b/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: - return inet_ntoa(pack('!i', ip_int)) - except struct_error: - # Backwards compatibility with old CTFd databases - return inet_ntoa(pack('!I', ip_int)) + '''Converts a saved IP address back into an integer/long''' + if ip_int < 4294967296: + # ipv4 + return inet_ntop(AF_INET, pack('!i', ip_int)) + else: + # ipv6 + return inet_ntop(AF_INET6, pack('!QQ', ip_int >> 64, ip_int & 0xffffffffffffffff)) db = SQLAlchemy() diff --git a/tests/test_utils.py b/tests/test_utils.py index e69de29b..dd2e5407 100644 --- a/tests/test_utils.py +++ b/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'