#6: added ability to pass custom dns servers to resolver

handler.py

@@ -270,3 +270,15 @@ async def data_handler(path: str) -> set:
         raw_data = read_file_to_list(path=path)
 
     return set(ip_converter(raw_data))
+
+
+def normalize_dns(dns_str: str) -> list:
+    # removes any whitespaces from string first. Further splits string into a list.
+    # comma is default delimiter.
+    dns_servers_list = (dns_str.replace(' ', '')).split(',')
+
+    valid_list = []
+    for dns_server_str in dns_servers_list:
+        if validators.ipv4(dns_server_str):
+            valid_list.append(dns_server_str)
+    return valid_list

rkn_exporter.py

@@ -13,7 +13,8 @@
     return_metrics, \
     fill_queue, \
     subnets_to_ips, \
-    data_handler
+    data_handler, \
+    normalize_dns
 
 # possibly it's good idea to use contextvars here
 data = 'rkn_computation_success 0'
@@ -52,6 +53,10 @@ def parse_args():
                         default=os.getenv("APP_THREADS_COUNT", 10),
                         type=int,
                         help='Threads count to parallelize computation. Is useful when DNS resolving is slow (default: 10)')
+    parser.add_argument('--dns',
+                        default=normalize_dns(os.getenv("APP_DNS", '8.8.8.8')),
+                        type=str,
+                        help='DNS servers (default: 8.8.8.8)')
     return parser.parse_args()
 
 
@@ -98,7 +103,7 @@ async def handler(self):
         blocked_ips_set = subnets_to_ips(blocked_subnets_set)
 
         # I'll add variables later
-        resolver = self.initialize_resolver(nameservers=['8.8.8.8'],
+        resolver = self.initialize_resolver(nameservers=self.args.dns,
                                             timeout=20,
                                             lifetime=20,
                                             retry_servfail=False)

tests/test_handler.py

@@ -17,7 +17,7 @@
                     return_domain_metrics, \
                     read_file_to_list, \
                     validate_domains, \
-                    fill_queue, subnet_to_ips, subnets_to_ips, get_data, ip_converter, data_handler
+                    fill_queue, subnet_to_ips, subnets_to_ips, get_data, ip_converter, data_handler, normalize_dns
 
 
 # taken from https://github.com/aio-libs/aiohttp/blob/master/tests/test_resolver.py
@@ -282,3 +282,10 @@ async def test_data_handler(mocker):
     mocker.patch('handler.read_file_to_list', return_value=await mock_awaitable_obj(raw_ips_list))
     blocked_subnets_set = await data_handler('./fake/path/to/file')
     assert blocked_subnets_set == valid_subnets_set
+
+
+def test_normalize_dns():
+    input_str = '  8.8.8.8 ,8.8.4.4,  8.8.1.1111'
+
+    valid_dns_list = normalize_dns(input_str)
+    assert valid_dns_list == ['8.8.8.8', '8.8.4.4']