diff --git a/routersploit/test/test_validators.py b/routersploit/test/test_validators.py index e7b2abe..8a2a658 100644 --- a/routersploit/test/test_validators.py +++ b/routersploit/test/test_validators.py @@ -73,13 +73,13 @@ class ValidatorsTest(RoutersploitTestCase): with self.assertRaises(OptionValidationError): validators.ipv4(address) - def test_ipv4_strip_scheme_1(self): + def test_address_strip_scheme_1(self): address = "http://127.0.0.1" - self.assertEqual(validators.ipv4(address), "127.0.0.1") + self.assertEqual(validators.address(address), "127.0.0.1") - def test_ipv4_strip_scheme_2(self): + def test_address_strip_scheme_2(self): address = "ftp://127.0.0.1" - self.assertEqual(validators.ipv4(address), "127.0.0.1") + self.assertEqual(validators.address(address), "127.0.0.1") if __name__ == '__main__': diff --git a/routersploit/validators.py b/routersploit/validators.py index da85ff2..10a7543 100644 --- a/routersploit/validators.py +++ b/routersploit/validators.py @@ -15,9 +15,12 @@ def url(address): return "http://{}".format(address) +def address(addr): + addr = urlparse.urlsplit(addr) + return addr.netloc or address.path + + def ipv4(address): - address = urlparse.urlsplit(address) - address = address.netloc or address.path try: socket.inet_pton(socket.AF_INET, address) except AttributeError: