From 0aa63fa2cf9ca10a7a6ccbe5b704f8ba1189117d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matteo=20=E2=84=B1an?= Date: Thu, 24 Sep 2020 22:44:15 +0200 Subject: [PATCH] Added timeout send/receive --- py-kms/pykms_Client.py | 99 +++++++++++++++++++++++++++--------------- py-kms/pykms_Misc.py | 9 ++++ py-kms/pykms_Server.py | 26 +++++------ 3 files changed, 84 insertions(+), 50 deletions(-) diff --git a/py-kms/pykms_Client.py b/py-kms/pykms_Client.py index ce1504d..bd6d3ec 100644 --- a/py-kms/pykms_Client.py +++ b/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 @@ -50,24 +50,28 @@ loggerclt = logging.getLogger('logclt') # 'help' string - 'default' value - 'dest' string. clt_options = { - 'ip' : {'help' : 'The IP address or hostname of the KMS server.', 'def' : "0.0.0.0", 'des' : "ip"}, - 'port' : {'help' : 'The port the KMS service is listening on. The default is \"1688\".', 'def' : 1688, 'des' : "port"}, - 'mode' : {'help' : 'Use this flag to manually specify a Microsoft product for testing the server. The default is \"Windows81\"', - 'def' : "Windows8.1", 'des' : "mode", - 'choi' : ["WindowsVista","Windows7","Windows8","Windows8.1","Windows10","Office2010","Office2013","Office2016","Office2019"]}, - 'cmid' : {'help' : 'Use this flag to manually specify a CMID to use. If no CMID is specified, a random CMID will be generated.', - '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 \ + 'ip' : {'help' : 'The IP address or hostname of the KMS server.', 'def' : "0.0.0.0", 'des' : "ip"}, + 'port' : {'help' : 'The port the KMS service is listening on. The default is \"1688\".', 'def' : 1688, 'des' : "port"}, + 'mode' : {'help' : 'Use this flag to manually specify a Microsoft product for testing the server. The default is \"Windows81\"', + 'def' : "Windows8.1", 'des' : "mode", + 'choi' : ["WindowsVista","Windows7","Windows8","Windows8.1","Windows10","Office2010","Office2013","Office2016","Office2019"]}, + 'cmid' : {'help' : 'Use this flag to manually specify a CMID to use. If no CMID is specified, a random CMID will be generated.', + '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", - 'choi' : ["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "MININFO"]}, - 'lfile' : {'help' : 'Use this option to set an output log file. The default is \"pykms_logclient.log\". \ + 'llevel' : {'help' : 'Use this option to set a log level. The default is \"ERROR\".', 'def' : "ERROR", 'des' : "loglevel", + 'choi' : ["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "MININFO"]}, + 'lfile' : {'help' : 'Use this option to set an output log file. The default is \"pykms_logclient.log\". \ Type \"STDOUT\" to view log info on stdout. Type \"FILESTDOUT\" to combine previous actions. \ Use \"STDOUTOFF\" to disable stdout messages. Use \"FILEOFF\" if you not want to create logfile.', - 'def' : os.path.join('.', 'pykms_logclient.log'), 'des' : "logfile"}, - 'lsize' : {'help' : 'Use this flag to set a maximum size (in MB) to the output log file. Deactivated by default.', 'def' : 0, 'des': "logsize"}, + 'def' : os.path.join('.', 'pykms_logclient.log'), 'des' : "logfile"}, + 'lsize' : {'help' : 'Use this flag to set a maximum size (in MB) to the output log file. Deactivated by default.', 'def' : 0, 'des': "logsize"}, } def client_options(): @@ -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,7 +148,11 @@ 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() @@ -164,17 +176,24 @@ def client_update(): clt_config['KMSClientAppID'] = appitem['Id'] clt_config['KMSClientKMSCountedID'] = kmsitem['Id'] break - -def client_create(): - loggerclt.info("Connecting to %s on port %d..." % (clt_config['ip'], clt_config['port'])) + +def client_connect(): + 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,20 +263,28 @@ 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): - if not with_gui: - # Parse options. - client_options() - - # Check options. - client_check() - # Update Config. - client_update() - # Create and run client. - client_create() - + try: + if not with_gui: + # Parse options. + client_options() + + # Check options. + client_check() + # Update Config. + client_update() + # Create and run client. + 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() requestDict['versionMinor'] = clt_config['KMSProtocolMinorVersion'] diff --git a/py-kms/pykms_Misc.py b/py-kms/pykms_Misc.py index 41a6a5c..7dbcabd 100644 --- a/py-kms/pykms_Misc.py +++ b/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) diff --git a/py-kms/pykms_Server.py b/py-kms/pykms_Server.py index 8dac940..8f30fb6 100755 --- a/py-kms/pykms_Server.py +++ b/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'] - list_opt = ['-c/--client-count', '-t0/--timeout-idle'] - + opts = [('clientcount', '-c/--client-count'), + ('timeoutidle', '-t0/--timeout-idle'), + ('timeoutsndrcv', '-t1/--timeout-sndrcv')] if serverthread.with_gui: - list_dest += ['activation', 'renewal'] - list_opt += ['-a/--activation-interval', '-r/--renewal-interval'] - - 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])) + opts += [('activation', '-a/--activation-interval'), + ('renewal', '-r/--renewal-interval')] + check_other(srv_config, opts, loggersrv, where = 'srv') # 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: