Harden ISO lookup handling

Author: Alcheri

__init__.py

@@ -32,8 +32,6 @@
 ISO: Convert alpha2 country codes to country name.
 """
 
-import sys
-import supybot
 from supybot import world
 
 # Use this for the version of this plugin.
@@ -60,7 +58,7 @@
 # reloaded when this plugin is reloaded.  Don't forget to import them as well!
 
 if world.testing:
-    from . import test
+    from . import test as test
 
 Class = plugin.Class
 configure = config.configure

config.py

@@ -28,25 +28,24 @@
 
 ###
 
-from supybot import conf, registry
+from supybot import conf
 
 try:
     from supybot.i18n import PluginInternationalization
 
     _ = PluginInternationalization("ISO")
-except:
+except ImportError:
     # Placeholder that allows to run the plugin on a bot
     # without the i18n module
-    _ = lambda x: x
+    def _(text):
+        return text
 
 
 def configure(advanced):
     # This will be called by supybot to configure this module.  advanced is
     # a bool that specifies whether the user identified themself as an advanced
     # user or not.  You should effect your configuration by manipulating the
     # registry as appropriate.
-    from supybot.questions import expect, anything, something, yn
-
     conf.registerPlugin("ISO", True)
 
 

plugin.py

@@ -29,7 +29,7 @@
 ###
 
 from supybot import callbacks
-from supybot.commands import *
+from supybot.commands import wrap
 
 try:
     from supybot.i18n import PluginInternationalization
@@ -38,14 +38,41 @@
 except ImportError:
     # Placeholder that allows to run the plugin on a bot
     # without the i18n module
-    _ = lambda x: x
+    def _(text):
+        return text
+
 
 try:
     from iso3166 import countries
 except ImportError as ie:
     raise ImportError(f"Cannot import module: {ie}")
 
 
+MAX_LOOKUP_LENGTH = 64
+
+
+def normalise_lookup(value):
+    """Normalise a country lookup value before querying iso3166."""
+    lookup = str(value or "").strip().lower()
+    if not lookup:
+        raise ValueError("Country code or name is required.")
+    if len(lookup) > MAX_LOOKUP_LENGTH:
+        raise ValueError("Country code or name is too long.")
+    return lookup
+
+
+def lookup_country(value):
+    """Return alpha-2 code and country name for a safe lookup value."""
+    lookup = normalise_lookup(value)
+    try:
+        country = countries.get(lookup)
+    except KeyError:
+        raise ValueError("Unknown country code or name.")
+    name = country[0]
+    alpha2 = country[1]
+    return alpha2, name
+
+
 class ISO(callbacks.Plugin):
     """Convert alpha2 country codes to country name."""
 
@@ -58,13 +85,10 @@ def country(self, irc, msg, args, code):
         Convert country name to alpha2 country codes.
         """
 
-        code = code.lower()
         try:
-            country = countries.get(code)
-            name = country[0]
-            alpha2 = country[1]
-        except KeyError as error:
-            raise callbacks.Error(f"{error} unknown country code.")
+            alpha2, name = lookup_country(code)
+        except ValueError as error:
+            raise callbacks.Error(str(error))
         irc.reply(f"{alpha2} {name}", prefixNick=False)
 
 

pyproject.toml

@@ -35,6 +35,15 @@ extend-exclude = '''
 )/
 '''
 
+[tool.mypy]
+python_version = "3.11"
+ignore_missing_imports = true
+explicit_package_bases = true
+exclude = "(__init__|test)\\.py"
+
+[tool.pytest.ini_options]
+python_files = ["test.py"]
+
 [tool.setuptools.package-dir]
 "limnoria_iso" = "."
 "limnoria_iso.local" = "local/"

test.py

@@ -28,11 +28,53 @@
 
 ###
 
-from supybot.test import *
+import unittest
 
+import supybot.test as supybot_test
 
-class ISOTestCase(PluginTestCase):
+supybot_test.PluginTestCase.__test__ = False
+
+try:
+    from . import plugin as iso_plugin
+except ImportError:  # pragma: no cover - allows direct unittest execution.
+    import plugin as iso_plugin
+
+
+class ISOTestCase(supybot_test.PluginTestCase):
+    __test__ = False
     plugins = ("ISO",)
 
 
+class ISOUnitTestCase(unittest.TestCase):
+    def test_lookup_country_accepts_alpha2_code(self):
+        self.assertEqual(iso_plugin.lookup_country("tr"), ("TR", "Türkiye"))
+
+    def test_lookup_country_accepts_country_name(self):
+        self.assertEqual(
+            iso_plugin.lookup_country("myanmar"), ("MM", "Myanmar")
+        )
+
+    def test_normalise_lookup_rejects_empty_input(self):
+        with self.assertRaisesRegex(
+            ValueError, "Country code or name is required."
+        ):
+            iso_plugin.normalise_lookup(" \n\t ")
+
+    def test_normalise_lookup_rejects_overlong_input(self):
+        with self.assertRaisesRegex(
+            ValueError, "Country code or name is too long."
+        ):
+            iso_plugin.normalise_lookup(
+                "x" * (iso_plugin.MAX_LOOKUP_LENGTH + 1)
+            )
+
+    def test_lookup_country_uses_generic_unknown_error(self):
+        with self.assertRaisesRegex(
+            ValueError, "Unknown country code or name."
+        ) as context:
+            iso_plugin.lookup_country("zz\nvery noisy input")
+
+        self.assertNotIn("zz", str(context.exception))
+
+
 # vim:set shiftwidth=4 tabstop=4 expandtab textwidth=79: