-
Notifications
You must be signed in to change notification settings - Fork 66
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: allow enforcing of the used IP version
Due to the need of whitelisting IP addresses, users might want to explicitly specify the used IP version
- Loading branch information
1 parent
b810c44
commit 3483566
Showing
4 changed files
with
105 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,50 @@ | ||
# -*- coding: utf-8 -*- | ||
import logging | ||
import socket | ||
|
||
import requests.packages.urllib3.util.connection as urllib3_cn | ||
|
||
|
||
def _allowed_gai_family_4(): | ||
"""https://github.com/urllib3/urllib3/blob/master/src/urllib3/util/connection.py""" | ||
return socket.AF_INET | ||
|
||
|
||
def _allowed_gai_family_6(): | ||
"""https://github.com/urllib3/urllib3/blob/master/src/urllib3/util/connection.py""" | ||
return socket.AF_INET6 | ||
|
||
|
||
def _allowed_gai_family_unspec(): | ||
"""https://github.com/urllib3/urllib3/blob/master/src/urllib3/util/connection.py""" | ||
return socket.AF_UNSPEC | ||
|
||
|
||
def enforce_ip_version(version=None): | ||
""" | ||
Tries to enforce either IPv4 or IPv6, depending on the version parameter | ||
:param version: The IP version to use | ||
:return: | ||
""" | ||
if version is None: | ||
urllib3_cn.allowed_gai_family = _allowed_gai_family_unspec | ||
return | ||
|
||
version = str(version) | ||
|
||
if version == "4": | ||
logging.info("Enforcing IPv4!") | ||
urllib3_cn.allowed_gai_family = _allowed_gai_family_4 | ||
|
||
elif version == "6": | ||
logging.info("Enforcing IPv6!") | ||
if not urllib3_cn.HAS_IPV6: | ||
raise Exception("Your system can't handle IPv6!") | ||
|
||
urllib3_cn.allowed_gai_family = _allowed_gai_family_6 | ||
|
||
elif version == "5": | ||
raise ValueError("Internet Stream Protocol? Really? There is no actual IPv5!") | ||
|
||
else: | ||
raise ValueError("No valid value specified for 'version' parameter! Please use either 4 or 6") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
import unittest | ||
from unittest.mock import Mock | ||
|
||
from pastepwn.util import enforce_ip_version | ||
|
||
import socket | ||
import requests.packages.urllib3.util.connection as urllib3_cn | ||
|
||
|
||
class TestNetwork(unittest.TestCase): | ||
"""Tests for the network utils - these tests might fail on CI""" | ||
|
||
def test_enforce_ip_version_4(self): | ||
enforce_ip_version(4) | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_INET) | ||
enforce_ip_version("4") | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_INET) | ||
|
||
def test_enforce_ip_version_5(self): | ||
with self.assertRaises(Exception): | ||
enforce_ip_version(5) | ||
|
||
with self.assertRaises(Exception): | ||
enforce_ip_version("5") | ||
|
||
def test_enforce_ip_version_6(self): | ||
enforce_ip_version(6) | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_INET6) | ||
enforce_ip_version("6") | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_INET6) | ||
|
||
def test_enforce_ip_version_None(self): | ||
enforce_ip_version(None) | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_UNSPEC) | ||
|
||
enforce_ip_version() | ||
self.assertEqual(urllib3_cn.allowed_gai_family(), socket.AF_UNSPEC) | ||
|
||
def test_enforce_ip_version_obj(self): | ||
mock = Mock() | ||
with self.assertRaises(ValueError): | ||
enforce_ip_version(mock) | ||
|
||
|
||
if __name__ == '__main__': | ||
unittest.main() |