class Script(object):
"""
Class for common CLI tool script
"""
def __init__(self, name=None, description=None, epilog=None, debug_flag=True):
self.name = os.path.basename(sys.argv[0])
setproctitle('%s %s' % (self.name, ' '.join(sys.argv[1:])))
signal.signal(signal.SIGINT, self.SIGINT)
reload(sys)
sys.setdefaultencoding('utf-8')
if name is None:
name = self.name
# Set to True to avoid any messages from self.message to be printed
self.silent = False
self.logger = Logger(self.name)
self.log = self.logger.default_stream
self.subcommand_parser = None
self.parser = argparse.ArgumentParser(
prog=name,
description=description,
formatter_class=argparse.RawTextHelpFormatter,
epilog=epilog,
add_help=True,
conflict_handler='resolve',
)
if debug_flag:
self.parser.add_argument('--debug', action='store_true', help='Show debug messages')
self.parser.add_argument('--insecure', action='store_false', help='No HTTPS certificate validation')
self.parser.add_argument('-B', '--browser',
choices=('chrome','chromium','firefox'),
help='Browser for cookie stealing'
)
def SIGINT(self, signum, frame):
"""
Parse SIGINT signal by quitting the program cleanly with exit code 1
"""
for t in filter(lambda t: t.name!='MainThread', threading.enumerate()):
t.stop()
for t in filter(lambda t: t.name!='MainThread', threading.enumerate()):
t.join()
self.exit(1)
def wait(self, poll_interval=1):
"""
Wait for running threads to finish.
Poll interval is time to wait between checks for threads
"""
while True:
active = filter(lambda t: t.name!='MainThread', threading.enumerate())
if not len(active):
break
self.log.debug('Waiting for %d threads' % len(active))
time.sleep(poll_interval)
def exit(self, value=0, message=None):
"""
Exit the script with given exit value.
If message is not None, it is printed on screen.
"""
if isinstance(value, bool):
if value:
value = 0
else:
value = 1
else:
try:
value = int(value)
if value < 0 or value > 255:
raise ValueError
except ValueError:
value = 1
if message is not None:
self.message(message)
for t in filter(lambda t: t.name!='MainThread', threading.enumerate()):
t.stop()
while True:
active = filter(lambda t: t.name!='MainThread', threading.enumerate())
if not len(active):
break
time.sleep(1)
sys.exit(value)
def message(self, message):
if self.silent:
return
sys.stdout.write('%s\n' % message)
def error(self, message):
sys.stderr.write('%s\n' % message)
def add_subcommand(self, command):
"""Add a subcommand parser instance
Register named subcommand parser to argument parser
Subcommand parser must be an instance of ScriptCommand class.
Example usage:
class ListCommand(ScriptCommand):
def run(self, args):
print 'Listing stuff'
parser.add_subcommand(ListCommand('list', 'List stuff from script'))
"""
if self.subcommand_parser is None:
self.subcommand_parser = self.parser.add_subparsers(
dest='command', help='Please select one command mode below',
title='Command modes'
)
self.subcommands = {}
if not isinstance(command, ScriptCommand):
raise ScriptError('Subcommand must be a ScriptCommand instance')
parser = self.subcommand_parser.add_parser(
command.name,
help=command.short_description,
description=command.description,
epilog=command.epilog,
formatter_class=argparse.RawTextHelpFormatter,
)
self.subcommands[command.name] = command
command.script = self
return parser
def usage_error(self, *args, **kwargs):
return self.parser.error(*args, **kwargs)
def add_argument(self, *args, **kwargs):
"""
Shortcut to add argument to main argumentparser instance
"""
self.parser.add_argument(*args, **kwargs)
def parse_args(self):
"""
Call parse_args for parser and check for default logging flags
"""
args = self.parser.parse_args()
if hasattr(args, 'debug') and getattr(args, 'debug'):
self.logger.set_level('DEBUG')
elif hasattr(args, 'quiet') and getattr(args, 'quiet'):
self.silent = True
elif hasattr(args, 'verbose') and getattr(args, 'verbose'):
self.logger.set_level('INFO')
if self.subcommand_parser is not None:
self.subcommands[args.command].run(args)
return args
def execute(self, args, dryrun=False):
"""
Default wrapper to execute given interactive shell command
with standard stdin, stdout and stderr
"""
if isinstance(args, basestring):
args = args.split()
if not isinstance(args, list):
raise ValueError('Execute arguments must be a list')
if dryrun:
self.log.debug('would execute: %s' % ' '.join(args))
return 0
p = Popen(args, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr)
p.wait()
return p.returncode
def check_output(self, args):
"""
Wrapper for subprocess.check_output to be executed in script context
"""
if isinstance(args, basestring):
args = [args]
try:
return check_output(args)
except IOError, (ecode, emsg):
raise ScriptError(emsg)
except OSError, (ecode, emsg):
raise ScriptError(emsg)
class Script(object):
"""
Class for common CLI tool script
"""
def __init__(self,
name=None,
description=None,
epilog=None,
debug_flag=True):
self.name = os.path.basename(sys.argv[0])
signal.signal(signal.SIGINT, self.SIGINT)
if has_setproctitle:
setproctitle('{0} {1}'.format(self.name, ' '.join(sys.argv[1:])))
if sys.version_info.major < 3:
reload(sys) # noqa
sys.setdefaultencoding('utf-8')
if name is None:
name = self.name
# Set to True to avoid any messages from self.message to be output
self.silent = False
self.logger = Logger(self.name)
self.log = self.logger.default_stream
self.parser = argparse.ArgumentParser(
prog=name,
description=description,
formatter_class=argparse.RawTextHelpFormatter,
epilog=epilog,
add_help=True,
conflict_handler='resolve',
)
if debug_flag:
self.parser.add_argument('--debug',
action='store_true',
help='Show debug messages')
self.subcommand_parser = None
def SIGINT(self, signum, frame):
"""
Parse SIGINT signal by quitting the program cleanly with exit code 1
"""
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
if hasattr(t, 'stop') and callable(t.stop):
t.stop()
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
t.join()
self.exit(1)
def wait(self, poll_interval=1):
"""
Wait for running threads to finish.
Poll interval is time to wait between checks for threads
"""
while True:
active = [
t for t in threading.enumerate() if t.name != 'MainThread'
]
if not len(active):
break
self.log.debug('Waiting for {0:d} threads'.format(len(active)))
time.sleep(poll_interval)
def exit(self, value=0, message=None):
"""
Exit the script with given exit value.
If message is not None, it is output to stdout
"""
if isinstance(value, bool):
if value:
value = 0
else:
value = 1
else:
try:
value = int(value)
if value < 0 or value > 255:
raise ValueError
except ValueError:
value = 1
if message is not None:
self.message(message)
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
if hasattr(t, 'stop') and callable(t.stop):
t.stop()
while True:
active = [
t for t in threading.enumerate() if t.name != 'MainThread'
]
if not len(active):
break
time.sleep(0.1)
sys.exit(value)
def message(self, message):
if self.silent:
return
sys.stdout.write('{0}\n'.format(message))
def error(self, message):
sys.stderr.write('{0}\n'.format(message))
def add_subcommand(self, command):
"""Add a subcommand parser instance
Register named subcommand parser to argument parser
Subcommand parser must be an instance of ScriptCommand class.
Example usage:
class ListCommand(ScriptCommand):
def run(self, args):
self.message('Listing stuff')
parser.add_subcommand(ListCommand('list', 'List stuff from script'))
"""
if self.subcommand_parser is None:
self.subcommand_parser = self.parser.add_subparsers(
dest='command',
help='Please select one command mode below',
title='Command modes')
self.subcommands = {}
if not isinstance(command, ScriptCommand):
raise ScriptError('Subcommand must be a ScriptCommand instance')
parser = self.subcommand_parser.add_parser(
command.name,
help=command.short_description,
description=command.description,
epilog=command.epilog,
formatter_class=argparse.RawTextHelpFormatter,
)
self.subcommands[command.name] = command
command.script = self
if callable(getattr(command, '__register_arguments__', None)):
command.__register_arguments__(parser)
return parser
def usage_error(self, *args, **kwargs):
return self.parser.error(*args, **kwargs)
def add_argument(self, *args, **kwargs):
"""
Shortcut to add argument to main argumentparser instance
"""
self.parser.add_argument(*args, **kwargs)
def __process_args__(self, args):
"""Process args
Process args from parse_*args CalledProcessError
"""
if getattr(args, 'debug', None):
self.logger.set_level('DEBUG')
elif getattr(args, 'quiet', None):
self.silent = True
elif getattr(args, 'verbose', None):
self.logger.set_level('INFO')
if self.subcommand_parser is not None and args.command is not None:
if hasattr(self.subcommands[args.command], 'parse_args'):
args = self.subcommands[args.command].parse_args(args)
self.subcommands[args.command].run(args)
return args
def parse_args(self):
"""
Call parse_args for parser and check for default logging flags
"""
return self.__process_args__(self.parser.parse_args())
def parse_known_args(self):
"""
Call parse_args for parser and check for default logging flags
"""
args, other_args = self.parser.parse_known_args()
args = self.__process_args__(args)
return args, other_args
def run(self):
"""
Run script, parsing arguments and running subcommands
This expects subcommands have been registered and exists if not
This simply runs self.parse_args() not expecting to do anything with
returned values since the subcommand is run.
"""
if self.subcommand_parser is None:
self.exit(1, 'Command defines no subcommands')
args = self.parse_args()
if args.command is None:
self.exit(1, 'No command selected')
def execute(self,
args,
stdin=sys.stdin,
stdout=sys.stdout,
stderr=sys.stderr,
dryrun=False):
"""
Default wrapper to execute given interactive shell command
with standard stdin, stdout and stderr
"""
if isinstance(args, str):
args = args.split()
if not isinstance(args, list):
raise ValueError('Execute arguments must be a list')
if dryrun:
self.log.debug('would execute: {0}'.format(' '.join(args)))
return 0
p = Popen(args, stdin=stdin, stdout=stdout, stderr=stderr)
p.wait()
return p.returncode
def check_output(self, args):
"""
Wrapper for subprocess.check_output to be executed in script context
"""
if isinstance(args, str):
args = [args]
try:
return check_output(args)
except IOError as e:
raise ScriptError(e)
except OSError as e:
raise ScriptError(e)
except CalledProcessError as e:
raise ScriptError(e)
class Script(object):
"""
Class for common CLI tool script
"""
def __init__(self, name=None, description=None, epilog=None, debug_flag=True):
self.name = os.path.basename(sys.argv[0])
signal.signal(signal.SIGINT, self.SIGINT)
if has_setproctitle:
setproctitle('{0} {1}'.format(self.name, ' '.join(sys.argv[1:])))
if sys.version_info.major < 3:
reload(sys) # noqa
sys.setdefaultencoding('utf-8')
if name is None:
name = self.name
# Set to True to avoid any messages from self.message to be output
self.silent = False
self.logger = Logger(self.name)
self.log = self.logger.default_stream
self.parser = argparse.ArgumentParser(
prog=name,
description=description,
formatter_class=argparse.RawTextHelpFormatter,
epilog=epilog,
add_help=True,
conflict_handler='resolve',
)
if debug_flag:
self.parser.add_argument('--debug', action='store_true', help='Show debug messages')
self.subcommand_parser = None
def SIGINT(self, signum, frame):
"""
Parse SIGINT signal by quitting the program cleanly with exit code 1
"""
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
if hasattr(t, 'stop') and callable(t.stop):
t.stop()
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
t.join()
self.exit(1)
def wait(self, poll_interval=1):
"""
Wait for running threads to finish.
Poll interval is time to wait between checks for threads
"""
while True:
active = [t for t in threading.enumerate() if t.name != 'MainThread']
if not len(active):
break
self.log.debug('Waiting for {0:d} threads'.format(len(active)))
time.sleep(poll_interval)
def exit(self, value=0, message=None):
"""
Exit the script with given exit value.
If message is not None, it is output to stdout
"""
if isinstance(value, bool):
if value:
value = 0
else:
value = 1
else:
try:
value = int(value)
if value < 0 or value > 255:
raise ValueError
except ValueError:
value = 1
if message is not None:
self.message(message)
for t in [t for t in threading.enumerate() if t.name != 'MainThread']:
if hasattr(t, 'stop') and callable(t.stop):
t.stop()
while True:
active = [t for t in threading.enumerate() if t.name != 'MainThread']
if not len(active):
break
time.sleep(0.1)
sys.exit(value)
def message(self, message):
if self.silent:
return
sys.stdout.write('{0}\n'.format(message))
def error(self, message):
sys.stderr.write('{0}\n'.format(message))
def add_subcommand(self, command):
"""Add a subcommand parser instance
Register named subcommand parser to argument parser
Subcommand parser must be an instance of ScriptCommand class.
Example usage:
class ListCommand(ScriptCommand):
def run(self, args):
self.message('Listing stuff')
parser.add_subcommand(ListCommand('list', 'List stuff from script'))
"""
if self.subcommand_parser is None:
self.subcommand_parser = self.parser.add_subparsers(
dest='command', help='Please select one command mode below',
title='Command modes'
)
self.subcommands = {}
if not isinstance(command, ScriptCommand):
raise ScriptError('Subcommand must be a ScriptCommand instance')
parser = self.subcommand_parser.add_parser(
command.name,
help=command.short_description,
description=command.description,
epilog=command.epilog,
formatter_class=argparse.RawTextHelpFormatter,
)
self.subcommands[command.name] = command
command.script = self
if callable(getattr(command, '__register_arguments__', None)):
command.__register_arguments__(parser)
return parser
def usage_error(self, *args, **kwargs):
return self.parser.error(*args, **kwargs)
def add_argument(self, *args, **kwargs):
"""
Shortcut to add argument to main argumentparser instance
"""
self.parser.add_argument(*args, **kwargs)
def __process_args__(self, args):
"""Process args
Process args from parse_*args CalledProcessError
"""
if getattr(args, 'debug', None):
self.logger.set_level('DEBUG')
elif getattr(args, 'quiet', None):
self.silent = True
elif getattr(args, 'verbose', None):
self.logger.set_level('INFO')
if self.subcommand_parser is not None and args.command is not None:
if hasattr(self.subcommands[args.command], 'parse_args'):
args = self.subcommands[args.command].parse_args(args)
self.subcommands[args.command].run(args)
return args
def parse_args(self):
"""
Call parse_args for parser and check for default logging flags
"""
return self.__process_args__(self.parser.parse_args())
def parse_known_args(self):
"""
Call parse_args for parser and check for default logging flags
"""
args, other_args = self.parser.parse_known_args()
args = self.__process_args__(args)
return args, other_args
def execute(self, args, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, dryrun=False):
"""
Default wrapper to execute given interactive shell command
with standard stdin, stdout and stderr
"""
if isinstance(args, str):
args = args.split()
if not isinstance(args, list):
raise ValueError('Execute arguments must be a list')
if dryrun:
self.log.debug('would execute: {0}'.format(' '.join(args)))
return 0
p = Popen(args, stdin=stdin, stdout=stdout, stderr=stderr)
p.wait()
return p.returncode
def check_output(self, args):
"""
Wrapper for subprocess.check_output to be executed in script context
"""
if isinstance(args, str):
args = [args]
try:
return check_output(args)
except IOError as e:
raise ScriptError(e)
except OSError as e:
raise ScriptError(e)
except CalledProcessError as e:
raise ScriptError(e)
