more error handling

dnstapir · Nov 15, 2024 · b3c937b · b3c937b
1 parent 42e5439
commit b3c937b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/dnstapir/dns/mozpsl.py b/dnstapir/dns/mozpsl.py
@@ -117,7 +117,12 @@ def load_psl(self, stream: io.StringIO) -> None:
 
     def coredomain(self, domain: str) -> tuple[str, str]:
         """Find ICANN and private name cut-off for domain"""
-        domain = domain.rstrip(".")
+        if not domain:
+            raise ValueError
+        try:
+            domain = domain.rstrip(".")
+        except AttributeError as exc:
+            raise ValueError from exc
         lbls = domain.split(".")
         lbls.reverse()
         c, p = self.trie.search(lbls)

diff --git a/tests/test_dns_mozpsl.py b/tests/test_dns_mozpsl.py
@@ -24,3 +24,14 @@ def test_mozpsl():
 
     with pytest.raises(KeyError):
         psl.coredomain("local.")
+
+    # IDN test
+    assert psl.coredomain("www.xn--mnchen-3ya.de.") == ("xn--mnchen-3ya.de", "")
+
+    # Edge cases
+    with pytest.raises(ValueError):
+        psl.coredomain("")
+    with pytest.raises(ValueError):
+        psl.coredomain(None)
+    with pytest.raises(KeyError):
+        psl.coredomain("invalid..domain.")