예제 #1
0
파일: shell.py 프로젝트: PanBen/jsontester
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)
예제 #2
0
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)
예제 #3
0
파일: shell.py 프로젝트: hile/systematic
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)