Added timeout send/receive

py-kms/pykms_Client.py

@@ -22,7 +22,7 @@ from pykms_RequestV5 import kmsRequestV5
 from pykms_RequestV6 import kmsRequestV6
 from pykms_RpcBase import rpcBase
 from pykms_DB2Dict import kmsDB2Dict
-from pykms_Misc import check_setup
+from pykms_Misc import check_setup, check_other
 from pykms_Misc import KmsParser, KmsParserException, KmsParserHelp
 from pykms_Misc import kms_parser_get, kms_parser_check_optionals, kms_parser_check_positionals
 from pykms_Format import justify, byterize, enco, deco, pretty_printer
@@ -59,6 +59,10 @@ clt_options = {
                       'def' : None, 'des' : "cmid"},
         'name'     : {'help' : 'Use this flag to manually specify an ASCII machine name to use. If no machine name is specified a random one \
 will be generated.', 'def' : None, 'des' : "machine"},
+        'time0'    : {'help' : 'Set the maximum time to wait for a connection attempt to KMS server to succeed. Default is no timeout.',
+                      'def' : None, 'des' : "timeoutidle"},
+        'time1'    : {'help' : 'Set the maximum time to wait for sending / receiving a request / response. Default is no timeout.',
+                      'def' : None, 'des' : "timeoutsndrcv"},
         'asyncmsg' : {'help' : 'Prints pretty / logging messages asynchronously. Deactivated by default.',
                       'def' : False, 'des' : "asyncmsg"},
         'llevel'   : {'help' : 'Use this option to set a log level. The default is \"ERROR\".', 'def' : "ERROR", 'des' : "loglevel",
@@ -82,6 +86,10 @@ def client_options():
                                    help = clt_options['cmid']['help'], type = str)
         client_parser.add_argument("-n", "--name", dest = clt_options['name']['des'] , default = clt_options['name']['def'],
                                    help = clt_options['name']['help'], type = str)
+        client_parser.add_argument("-t0", "--timeout-idle", action = "store", dest = clt_options['time0']['des'], default = clt_options['time0']['def'],
+                                   help = clt_options['time0']['help'], type = str)
+        client_parser.add_argument("-t1", "--timeout-sndrcv", action = "store", dest = clt_options['time1']['des'], default = clt_options['time1']['def'],
+                                   help = clt_options['time1']['help'], type = str)
         client_parser.add_argument("-y", "--async-msg", action = "store_true", dest = clt_options['asyncmsg']['des'],
                                    default = clt_options['asyncmsg']['def'], help = clt_options['asyncmsg']['help'])
         client_parser.add_argument("-V", "--loglevel", dest = clt_options['llevel']['des'], action = "store",
@@ -140,6 +148,10 @@ def client_check():
                         
         clt_config['call_id'] = 1
 
+        # Check other specific client options.
+        opts = [('timeoutidle', '-t0/--timeout-idle'),
+                ('timeoutsndrcv', '-t1/--timeout-sndrcv')]
+        check_other(clt_config, opts, loggerclt, where = 'clt')
 
 def client_update():
         kmsdb = kmsDB2Dict()
@@ -165,16 +177,23 @@ def client_update():
                                                 clt_config['KMSClientKMSCountedID'] = kmsitem['Id']
                                                 break
 
-def client_create():
+def client_connect():
-        loggerclt.info("Connecting to %s on port %d..." % (clt_config['ip'], clt_config['port']))
+        loggerclt.info("Connecting to %s on port %d" % (clt_config['ip'], clt_config['port']))
         try:
-                clt_sock = socket.create_connection((clt_config['ip'], clt_config['port']))
+                clt_sock = socket.create_connection((clt_config['ip'], clt_config['port']), timeout = clt_config['timeoutidle'])
                 loggerclt.info("Connection successful !")
+                clt_sock.settimeout(clt_config['timeoutsndrcv'])
+        except socket.timeout:
+                pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
+                               put_text = "{reverse}{red}{bold}Client connection timed out. Exiting...{end}")
         except (socket.gaierror, socket.error) as e:
                 pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
                                put_text = "{reverse}{red}{bold}Connection failed '%s:%d': %s. Exiting...{end}" %(clt_config['ip'],
                                                                                                                  clt_config['port'],
                                                                                                                  str(e)))
+        return clt_sock
+
+def client_create(clt_sock):
         binder = pykms_RpcBind.handler(None, clt_config)
         RPC_Bind = enco(str(binder.generateRequest()), 'latin-1')
 
@@ -184,16 +203,16 @@ def client_create():
                 clt_sock.send(RPC_Bind)
         except socket.error as e:
                 pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
-                               put_text = "{reverse}{red}{bold}While sending: %s{end}" %str(e))
+                               put_text = "{reverse}{red}{bold}While sending: %s. Exiting...{end}" %str(e))
         try:
                 bindResponse = clt_sock.recv(1024)
                 if bindResponse == '' or not bindResponse:
                         pretty_printer(log_obj = loggerclt.warning, to_exit = True, where = "clt",
-                                       put_text = "{reverse}{yellow}{bold}No data received.{end}")
+                                       put_text = "{reverse}{yellow}{bold}No data received. Exiting...{end}")
                 pretty_printer(num_text = [-4, 7], where = "clt")
         except socket.error as e:
                 pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
-                               put_text = "{reverse}{red}{bold}While receiving: %s{end}" %str(e))
+                               put_text = "{reverse}{red}{bold}While receiving: %s. Exiting...{end}" %str(e))
 
         packetType = MSRPCHeader(bindResponse)['type']
         if packetType == rpcBase.packetType['bindAck']:
@@ -209,13 +228,13 @@ def client_create():
                         clt_sock.send(RPC_Actv)
                 except socket.error as e:
                         pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
-                                       put_text = "{reverse}{red}{bold}While sending: %s{end}" %str(e))
+                                       put_text = "{reverse}{red}{bold}While sending: %s. Exiting...{end}" %str(e))
                 try:
                         response = clt_sock.recv(1024)
                         pretty_printer(num_text = [-4, 20], where = "clt")
                 except socket.error as e:
                         pretty_printer(log_obj = loggerclt.error, to_exit = True, where = "clt",
-                                       put_text = "{reverse}{red}{bold}While receiving: %s{end}" %str(e))
+                                       put_text = "{reverse}{red}{bold}While receiving: %s. Exiting...{end}" %str(e))
 
                 loggerclt.debug("Response: \n%s\n" % justify(deco(binascii.b2a_hex(response), 'latin-1')))
                 parsed = MSRPCRespHeader(response)
@@ -244,9 +263,10 @@ def client_create():
                 sys.exit(0)
         else:
                 pretty_printer(log_obj = loggerclt.warning, to_exit = True, where = "clt",
-                               put_text = "{reverse}{magenta}{bold}Something went wrong.{end}")
+                               put_text = "{reverse}{magenta}{bold}Something went wrong. Exiting...{end}")
 
 def clt_main(with_gui = False):
+        try:
                 if not with_gui:
                         # Parse options.
                         client_options()
@@ -256,7 +276,14 @@ def clt_main(with_gui = False):
                 # Update Config.
                 client_update()
                 # Create and run client.
-        client_create()
+                clt_sock = client_connect()
+                client_create(clt_sock)
+        except (KeyboardInterrupt, SystemExit):
+                try:
+                        clt_sock.shutdown(socket.SHUT_RDWR)
+                        clt_sock.close()
+                except:
+                        pass
 
 def createKmsRequestBase():
         requestDict = kmsBase.kmsRequestStruct()

py-kms/pykms_Misc.py

@@ -552,6 +552,15 @@ def check_setup(config, options, logger, where):
                 pretty_printer(log_obj = logger.error, where = where, to_exit = True,
                                put_text = "{reverse}{red}{bold}Port number '%s' is invalid. Enter between 1 - 65535. Exiting...{end}" %config['port'])
 
+def check_other(config, options, logger, where):
+        for dest, stropt in options:
+                try:
+                        config[dest] = int(config[dest])
+                except:
+                        if config[dest] is not None:
+                                pretty_printer(log_obj = logger.error, where = where, to_exit = True,
+                                               put_text = "{reverse}{red}{bold}argument `%s`: invalid with: '%s'. Exiting...{end}" %(stropt, config[dest]))
+
 #------------------------------------------------------------------------------------------------------------------------------------------------------------
 
 # http://joshpoley.blogspot.com/2011/09/hresults-user-0x004.html  (slerror.h)

py-kms/pykms_Server.py

@@ -18,7 +18,7 @@ from time import monotonic as time
 import pykms_RpcBind, pykms_RpcRequest
 from pykms_RpcBase import rpcBase
 from pykms_Dcerpc import MSRPCHeader
-from pykms_Misc import check_setup, check_lcid, check_dir
+from pykms_Misc import check_setup, check_lcid, check_dir, check_other
 from pykms_Misc import KmsParser, KmsParserException, KmsParserHelp
 from pykms_Misc import kms_parser_get, kms_parser_check_optionals, kms_parser_check_positionals, kms_parser_check_connect
 from pykms_Format import enco, deco, pretty_printer, justify
@@ -202,6 +202,8 @@ The default is \"364F463A8863D35F\" or type \"RANDOM\" to auto generate the HWID
                         'def' : "364F463A8863D35F", 'des' : "hwid"},
         'time0'      : {'help' : 'Maximum inactivity time (in seconds) after which the connection with the client is closed. If \"None\" (default) serve forever.',
                         'def' : None, 'des' : "timeoutidle"},
+        'time1'      : {'help' : 'Set the maximum time to wait for sending / receiving a request / response. Default is no timeout.',
+                        'def' : None, 'des' : "timeoutsndrcv"},
         'asyncmsg'   : {'help' : 'Prints pretty / logging messages asynchronously. Deactivated by default.',
                         'def' : False, 'des' : "asyncmsg"},
         'llevel'     : {'help' : 'Use this option to set a log level. The default is \"ERROR\".', 'def' : "ERROR", 'des' : "loglevel",
@@ -238,6 +240,8 @@ def server_options():
                                    help = srv_options['hwid']['help'], type = str)
         server_parser.add_argument("-t0", "--timeout-idle", action = "store", dest = srv_options['time0']['des'], default = srv_options['time0']['def'],
                                    help = srv_options['time0']['help'], type = str)
+        server_parser.add_argument("-t1", "--timeout-sndrcv", action = "store", dest = srv_options['time1']['des'], default = srv_options['time1']['def'],
+                                   help = srv_options['time1']['help'], type = str)
         server_parser.add_argument("-y", "--async-msg", action = "store_true", dest = srv_options['asyncmsg']['des'],
                                    default = srv_options['asyncmsg']['def'], help = srv_options['asyncmsg']['help'])
         server_parser.add_argument("-V", "--loglevel", action = "store", dest = srv_options['llevel']['des'], choices = srv_options['llevel']['choi'],
@@ -478,20 +482,13 @@ def server_check():
                         srv_config['sqlite'] = False
 
         # Check other specific server options.
-        list_dest = ['clientcount', 'timeoutidle']
+        opts = [('clientcount', '-c/--client-count'),
-        list_opt = ['-c/--client-count', '-t0/--timeout-idle']
+                ('timeoutidle', '-t0/--timeout-idle'),
-
+                ('timeoutsndrcv', '-t1/--timeout-sndrcv')]
         if serverthread.with_gui:
-                list_dest += ['activation', 'renewal']
+                opts += [('activation', '-a/--activation-interval'),
-                list_opt += ['-a/--activation-interval', '-r/--renewal-interval']
+                         ('renewal', '-r/--renewal-interval')]
-
+        check_other(srv_config, opts, loggersrv, where = 'srv')
-        for dest, opt in zip(list_dest, list_opt):
-                try:
-                        srv_config[dest] = int(srv_config[dest])
-                except:
-                        if srv_config[dest] is not None:
-                                pretty_printer(log_obj = loggersrv.error, to_exit = True,
-                                               put_text = "{reverse}{red}{bold}argument `%s`: invalid with: '%s'. Exiting...{end}" %(opt, srv_config[dest]))
 
         # Check further addresses / ports.
         if 'listen' in srv_config:
@@ -602,6 +599,7 @@ class kmsServerHandler(socketserver.BaseRequestHandler):
                 srv_config['raddr'] = self.client_address
 
         def handle(self):
+                self.request.settimeout(srv_config['timeoutsndrcv'])
                 while True:
                         # self.request is the TCP socket connected to the client
                         try: