diff --git a/src/mvt/common/url.py b/src/mvt/common/url.py index 426d64bd..d432e3ed 100644 --- a/src/mvt/common/url.py +++ b/src/mvt/common/url.py @@ -349,7 +349,9 @@ class URL: return tld_obj if tld_obj is None: return "" - return tld_obj.parsed_url.netloc.lower().lstrip("www.") + # removeprefix, not lstrip: lstrip takes a set of characters, so it ate the + # leading "w"s and dots of any domain ("web.evil.com" -> "eb.evil.com"). + return tld_obj.parsed_url.netloc.lower().removeprefix("www.") def get_top_level(self) -> str: """Get only the top-level domain from a URL. diff --git a/tests/common/test_url.py b/tests/common/test_url.py index ce0fb635..52bfb5e3 100644 --- a/tests/common/test_url.py +++ b/tests/common/test_url.py @@ -22,3 +22,18 @@ def test_google_maps_url_is_not_shortened(url): def test_other_google_short_url_is_shortened(): assert URL("https://goo.gl/example").check_if_shortened() is True + + +@pytest.mark.parametrize( + "url, domain", + [ + ("https://www.example.com/path", "example.com"), + # Only the whole "www." prefix comes off, not any leading "w" or "." character. + ("https://web.example.com", "web.example.com"), + ("https://wow.com", "wow.com"), + ("https://wired.com", "wired.com"), + ("https://www.wow.com", "wow.com"), + ], +) +def test_get_domain_strips_only_a_whole_www_prefix(url, domain): + assert URL(url).domain == domain