From 11d06e2be75875a22bec20c6dbf03ecf6d06e883 Mon Sep 17 00:00:00 2001
From: Nick Gheorghita <nickgheorghita@gmail.com>
Date: Mon, 14 Jan 2019 17:59:15 +0100
Subject: [PATCH] Enable additional tlds via environment variables for ens

---
 ens/constants.py                |  2 +-
 ens/utils.py                    | 20 ++++++++++++++++----
 tests/ens/test_setup_address.py |  2 +-
 tests/ens/test_setup_name.py    |  2 +-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/ens/constants.py b/ens/constants.py
index 2f0971a362..fb7e26a9cb 100644
--- a/ens/constants.py
+++ b/ens/constants.py
@@ -9,6 +9,6 @@
 
 MIN_ETH_LABEL_LENGTH = 7
 
-RECOGNIZED_TLDS = ['eth', 'reverse', 'test', 'luxe', 'xyz']
+DEFAULT_RECOGNIZED_TLDS = ['eth', 'reverse', 'test', 'luxe', 'xyz']
 
 REVERSE_REGISTRAR_DOMAIN = 'addr.reverse'
diff --git a/ens/utils.py b/ens/utils.py
index 91f18a513c..2b28dc4c3c 100644
--- a/ens/utils.py
+++ b/ens/utils.py
@@ -1,7 +1,7 @@
-
 import copy
 import datetime
 import functools
+import os
 
 from eth_utils import (
     is_same_address,
@@ -14,9 +14,9 @@
     ACCEPTABLE_STALE_HOURS,
     AUCTION_START_GAS_CONSTANT,
     AUCTION_START_GAS_MARGINAL,
+    DEFAULT_RECOGNIZED_TLDS,
     EMPTY_SHA3_BYTES,
     MIN_ETH_LABEL_LENGTH,
-    RECOGNIZED_TLDS,
     REVERSE_REGISTRAR_DOMAIN,
 )
 from ens.exceptions import (
@@ -114,13 +114,25 @@ def label_to_name(label, recognized_tlds):
     pieces = label.split('.')
     if pieces[-1] not in recognized_tlds:
         raise InvalidTLD(
-            f"Label does not have a recognized TLD. TLD must match one of: {recognized_tlds}."
+            f"The label: {label} has an unsupported TLD of {pieces[-1]}. "
+            f"ENS.py by default supports the following TLDs: {recognized_tlds}. "
+            "If you'd like to use an unsupported TLD, please set the environment variable: "
+            "'ENS_RECOGNIZED_TLDS' to a string of desired TLDs separated by a colon (:)."
         )
     return '.'.join(pieces)
 
 
 def dot_eth_name(label):
-    return label_to_name(label, RECOGNIZED_TLDS)
+    recognized_tlds = get_recognized_tlds()
+    return label_to_name(label, recognized_tlds)
+
+
+def get_recognized_tlds():
+    if 'ENS_RECOGNIZED_TLDS' in os.environ:
+        override_tlds = os.environ['ENS_RECOGNIZED_TLDS'].split(':')
+        return set(DEFAULT_RECOGNIZED_TLDS + override_tlds)
+    else:
+        return DEFAULT_RECOGNIZED_TLDS
 
 
 def name_to_label(name, registrar):
diff --git a/tests/ens/test_setup_address.py b/tests/ens/test_setup_address.py
index 547a032335..e29b544263 100644
--- a/tests/ens/test_setup_address.py
+++ b/tests/ens/test_setup_address.py
@@ -92,7 +92,7 @@ def test_set_address(ens, name, full_name, namehash_hex, TEST_ADDRESS):
     ),
 )
 def test_set_address_raises_exception_with_invalid_or_missing_tld(ens, name, TEST_ADDRESS):
-    with pytest.raises(InvalidTLD, match="Label does not have a recognized TLD."):
+    with pytest.raises(InvalidTLD, match="ENS.py by default supports the following TLDs"):
         ens.setup_address(name, TEST_ADDRESS)
 
 
diff --git a/tests/ens/test_setup_name.py b/tests/ens/test_setup_name.py
index 48fc50e7dd..4532c0b7f3 100644
--- a/tests/ens/test_setup_name.py
+++ b/tests/ens/test_setup_name.py
@@ -89,7 +89,7 @@ def test_setup_name(ens, name, normalized_name, namehash_hex):
 def test_cannot_setup_name_with_missing_or_invalid_tld(ens, name):
     address = ens.web3.eth.accounts[3]
     assert not ens.name(address)
-    with pytest.raises(InvalidTLD, match="Label does not have a recognized TLD."):
+    with pytest.raises(InvalidTLD, match="ENS.py by default supports the following TLDs"):
         ens.setup_name(name, address)