mirror of
https://github.com/SystemRage/py-kms.git
synced 2024-11-25 17:55:37 +01:00
Added timeout send/receive
This commit is contained in:
Matteo ℱan
parent
990cd5e48f
commit
0aa63fa2cf
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user