Added timeout send/receive

This commit is contained in:
Matteo ℱan 2020-09-24 22:44:15 +02:00
parent 990cd5e48f
commit 0aa63fa2cf
No known key found for this signature in database
GPG Key ID: 3C30A05BC133D9B6
3 changed files with 84 additions and 50 deletions

View File

@ -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():
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,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()

View File

@ -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)

View File

@ -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: