From c74a64cc6722cfcdb72ee347cfe859a502b583d6 Mon Sep 17 00:00:00 2001
From: edgd1er <edgd1er@hotmail.com>
Date: Thu, 14 Apr 2022 07:44:08 +0200
Subject: [PATCH] client: handle dns exceptions #51

---
 docs/Usage.md          |  4 ++--
 py-kms/pykms_Client.py | 16 +++++++++++-----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/docs/Usage.md b/docs/Usage.md
index ab69d66..2c26824 100644
--- a/docs/Usage.md
+++ b/docs/Usage.md
@@ -143,9 +143,9 @@ user@host ~/path/to/folder/py-kms $ python3 pykms_Server.py -V DEBUG
 user@host ~/path/to/folder/py-kms $ python3 pykms_Client.py -V DEBUG
 ```
 
-If you wish to get KMS server from DNS server: (ie perform a DNS resolution on _vlmcs._tcp.domain.tld). If ever there are several answers, only the first one is selected.
+If you wish to get KMS server from DNS server: (ie perform a DNS resolution on _vlmcs._tcp.domain.tld, if ever there are several answers, only the first one is selected.)
 ```
-user@host ~/path/to/folder/py-kms $ python3 pykms_Server.py -V DEBUG -F STDOUT -D contoso.com
+user@host ~/path/to/folder/py-kms $ python3 pykms_Client.py -V DEBUG -F STDOUT -D contoso.com
 user@host ~/path/to/folder/py-kms $ python3 pykms_Client.py -V DEBUG -F STDOUT -D contoso.com
 ```
 
diff --git a/py-kms/pykms_Client.py b/py-kms/pykms_Client.py
index b203a6f..109f83a 100644
--- a/py-kms/pykms_Client.py
+++ b/py-kms/pykms_Client.py
@@ -198,11 +198,17 @@ def client_connect():
 
         if clt_config['discovery'] is not None:
           loggerclt.info(f'Using Domain: {clt_config["discovery"]}')
-          r = dns.resolver.query('_vlmcs._tcp.' + clt_config['discovery'], dns.rdatatype.SRV)
-          for a in r:
-            loggerclt.debug(f'answer KMS server: {a.target} , port: {a.port}')
-          clt_config['ip'] = socket.gethostbyname(r[0].target.to_text())
-          clt_config['port'] = r[0].port
+          r= None
+          try:
+            r = dns.resolver.resolve('_vlmcs._tcp.' + clt_config['discovery'], dns.rdatatype.SRV)
+            for a in r:
+              loggerclt.debug(f'answer KMS server: {a.target} , port: {a.port}')
+            clt_config['ip'] = socket.gethostbyname(r[0].target.to_text())
+            clt_config['port'] = r[0].port
+          except (dns.exception.Timeout, dns.resolver.NXDOMAIN) as e:
+                pretty_printer(log_obj = loggerclt.warning,
+                           put_text = "{reverse}{red}{bold}Cannot resolve '%s'. Error: '%s'...{end}" %(clt_config['discovery'],
+                                                                                                             str(e)))
 
         loggerclt.info("Connecting to %s on port %d" % (clt_config['ip'], clt_config['port']))
         try: