"Fossies" - the Fresh Open Source Software Archive

Member "polysh-polysh-0.13/polysh/main.py" (11 May 2020, 9715 Bytes) of package /linux/privat/polysh-polysh-0.13.tar.gz:

    1 """Polysh - Main Utilities
    2 
    3 Copyright (c) 2006 Guillaume Chazarain <guichaz@gmail.com>
    4 Copyright (c) 2018 InnoGames GmbH
    5 """
    6 # This program is free software: you can redistribute it and/or modify
    7 # it under the terms of the GNU General Public License as published by
    8 # the Free Software Foundation, either version 2 of the License, or
    9 # (at your option) any later version.
   10 #
   11 # This program is distributed in the hope that it will be useful,
   12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
   13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   14 # GNU General Public License for more details.
   15 #
   16 # You should have received a copy of the GNU General Public License
   17 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
   18 
   19 import asyncore
   20 import atexit
   21 import getpass
   22 import locale
   23 import argparse
   24 import os
   25 import signal
   26 import sys
   27 import termios
   28 import readline
   29 import resource
   30 from typing import Callable, List
   31 
   32 from polysh import remote_dispatcher
   33 from polysh import dispatchers
   34 from polysh import stdin
   35 from polysh.console import console_output
   36 from polysh.host_syntax import expand_syntax
   37 from polysh import control_commands
   38 from polysh import VERSION
   39 
   40 
   41 def kill_all() -> None:
   42     """When polysh quits, we kill all the remote shells we started"""
   43     for i in dispatchers.all_instances():
   44         try:
   45             os.kill(-i.pid, signal.SIGKILL)
   46         except OSError:
   47             # The process was already dead, no problem
   48             pass
   49 
   50 
   51 def parse_cmdline() -> argparse.Namespace:
   52     usage = '%s [OPTIONS] HOSTS...\n' % (sys.argv[0]) + \
   53             'Control commands are prefixed by ":".'
   54     parser = argparse.ArgumentParser(usage)
   55     parser.add_argument(
   56         '--hosts-file', type=str, action='append',
   57         dest='hosts_filenames', metavar='FILE', default=[],
   58         help='read hostnames from given file, one per line')
   59     parser.add_argument(
   60         '--command', type=str, dest='command', default=None,
   61         help='command to execute on the remote shells',
   62         metavar='CMD')
   63     def_ssh = 'exec ssh -oLogLevel=Quiet -t %(host)s %(port)s'
   64     parser.add_argument(
   65         '--ssh', type=str, dest='ssh', default=def_ssh,
   66         metavar='SSH', help='ssh command to use [%s]' % def_ssh)
   67     parser.add_argument(
   68         '--user', type=str, dest='user', default=None,
   69         help='remote user to log in as', metavar='USER')
   70     parser.add_argument(
   71         '--no-color', action='store_true', dest='disable_color',
   72         help='disable colored hostnames [enabled]')
   73     parser.add_argument(
   74         '--password-file', type=str, dest='password_file',
   75         default=None, metavar='FILE',
   76         help='read a password from the specified file. - is the tty.')
   77     parser.add_argument(
   78         '--log-file', type=str, dest='log_file',
   79         help='file to log each machine conversation [none]')
   80     parser.add_argument(
   81         '--abort-errors', action='store_true', dest='abort_error',
   82         help='abort if some shell fails to initialize [ignore]')
   83     parser.add_argument(
   84         '--debug', action='store_true', dest='debug',
   85         help='print debugging information')
   86     parser.add_argument(
   87         '--profile', action='store_true', dest='profile',
   88         default=False)
   89     parser.add_argument('host_names', nargs='*')
   90     args = parser.parse_args()
   91 
   92     for filename in args.hosts_filenames:
   93         try:
   94             hosts_file = open(filename, 'r')
   95             for line in hosts_file.readlines():
   96                 if '#' in line:
   97                     line = line[:line.index('#')]
   98                 line = line.strip()
   99                 if line:
  100                     args.host_names.append(line)
  101             hosts_file.close()
  102         except IOError as e:
  103             parser.error(str(e))
  104 
  105     if args.log_file:
  106         try:
  107             args.log_file = open(args.log_file, 'a')
  108         except IOError as e:
  109             print(e)
  110             sys.exit(1)
  111 
  112     if not args.host_names:
  113         parser.error('no hosts given')
  114 
  115     if args.password_file == '-':
  116         args.password = getpass.getpass()
  117     elif args.password_file is not None:
  118         password_file = open(args.password_file, 'r')
  119         args.password = password_file.readline().rstrip('\n')
  120     else:
  121         args.password = None
  122 
  123     return args
  124 
  125 
  126 def find_non_interactive_command(command: str) -> str:
  127     if sys.stdin.isatty():
  128         return command
  129 
  130     stdin = sys.stdin.read()
  131     if stdin and command:
  132         print(
  133             '--command and reading from stdin are incompatible',
  134             file=sys.stderr,
  135         )
  136         sys.exit(1)
  137     if stdin and not stdin.endswith('\n'):
  138         stdin += '\n'
  139     return command or stdin
  140 
  141 
  142 def init_history(histfile: str) -> None:
  143     if hasattr(readline, "read_history_file"):
  144         try:
  145             readline.read_history_file(histfile)
  146         except IOError:
  147             pass
  148 
  149 
  150 def save_history(histfile: str) -> None:
  151     readline.set_history_length(1000)
  152     readline.write_history_file(histfile)
  153 
  154 
  155 def loop(interactive: bool) -> None:
  156     histfile = os.path.expanduser("~/.polysh_history")
  157     init_history(histfile)
  158     next_signal = None
  159     last_status = None
  160     while True:
  161         try:
  162             if next_signal:
  163                 current_signal = next_signal
  164                 next_signal = None
  165                 sig2chr = {signal.SIGINT: 'C', signal.SIGTSTP: 'Z'}
  166                 ctrl = sig2chr[current_signal]
  167                 remote_dispatcher.log('> ^{}\n'.format(ctrl).encode())
  168                 control_commands.do_send_ctrl(ctrl)
  169                 console_output(b'')
  170                 stdin.the_stdin_thread.prepend_text = None
  171             while dispatchers.count_awaited_processes()[0] and \
  172                     remote_dispatcher.main_loop_iteration(timeout=0.2):
  173                 pass
  174             # Now it's quiet
  175             for r in dispatchers.all_instances():
  176                 r.print_unfinished_line()
  177             current_status = dispatchers.count_awaited_processes()
  178             if current_status != last_status:
  179                 console_output(b'')
  180             if remote_dispatcher.options.interactive:
  181                 stdin.the_stdin_thread.want_raw_input()
  182             last_status = current_status
  183             if dispatchers.all_terminated():
  184                 # Clear the prompt
  185                 console_output(b'')
  186                 raise asyncore.ExitNow(remote_dispatcher.options.exit_code)
  187             if not next_signal:
  188                 # possible race here with the signal handler
  189                 remote_dispatcher.main_loop_iteration()
  190         except KeyboardInterrupt:
  191             if interactive:
  192                 next_signal = signal.SIGINT
  193             else:
  194                 kill_all()
  195                 os.kill(0, signal.SIGINT)
  196         except asyncore.ExitNow as e:
  197             console_output(b'')
  198             save_history(histfile)
  199             sys.exit(e.args[0])
  200 
  201 
  202 def _profile(continuation: Callable) -> None:
  203     prof_file = 'polysh.prof'
  204     import cProfile
  205     import pstats
  206     print('Profiling using cProfile')
  207     cProfile.runctx('continuation()', globals(), locals(), prof_file)
  208     stats = pstats.Stats(prof_file)
  209     stats.strip_dirs()
  210     stats.sort_stats('time', 'calls')
  211     stats.print_stats(50)
  212     stats.print_callees(50)
  213     os.remove(prof_file)
  214 
  215 
  216 def restore_tty_on_exit() -> None:
  217     fd = sys.stdin.fileno()
  218     old = termios.tcgetattr(fd)
  219     atexit.register(lambda: termios.tcsetattr(fd, termios.TCSADRAIN, old))
  220 
  221 
  222 def run() -> None:
  223     """Launch polysh"""
  224     locale.setlocale(locale.LC_ALL, '')
  225     atexit.register(kill_all)
  226     signal.signal(signal.SIGPIPE, signal.SIG_DFL)
  227 
  228     args = parse_cmdline()
  229 
  230     args.command = find_non_interactive_command(args.command)
  231     args.exit_code = 0
  232     args.interactive = (
  233         not args.command
  234         and sys.stdin.isatty()
  235         and sys.stdout.isatty())
  236     if args.interactive:
  237         restore_tty_on_exit()
  238 
  239     remote_dispatcher.options = args
  240 
  241     hosts = []  # type: List[str]
  242     for host in args.host_names:
  243         hosts.extend(expand_syntax(host))
  244 
  245     try:
  246         # stdin, stdout, stderr for polysh and each ssh connection
  247         new_soft = 3 + len(hosts) * 3
  248         old_soft, old_hard = resource.getrlimit(resource.RLIMIT_NOFILE)
  249         if new_soft > old_soft:
  250             # We are allowed to change the soft limit as we please but must be
  251             # root to change the hard limit.
  252             new_hard = max(new_soft, old_hard)
  253             resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, new_hard))
  254     except OSError as e:
  255         print(
  256             'Failed to change RLIMIT_NOFILE from soft={} hard={} to soft={} '
  257             'hard={}: {}'.format(old_soft, old_hard, new_soft, new_hard, e),
  258             file=sys.stderr,
  259         )
  260         sys.exit(1)
  261 
  262     dispatchers.create_remote_dispatchers(hosts)
  263 
  264     signal.signal(signal.SIGWINCH, lambda signum, frame:
  265                   dispatchers.update_terminal_size())
  266 
  267     stdin.the_stdin_thread = stdin.StdinThread(args.interactive)
  268 
  269     if args.profile:
  270         def safe_loop() -> None:
  271             try:
  272                 loop(args.interactive)
  273             except BaseException:
  274                 pass
  275         _profile(safe_loop)
  276     else:
  277         loop(args.interactive)
  278 
  279 
  280 def main():
  281     """Wrapper around run() to setup sentry"""
  282 
  283     sentry_dsn = os.environ.get('POLYSH_SENTRY_DSN')
  284 
  285     if sentry_dsn:
  286         from raven import Client
  287         client = Client(
  288             dsn=sentry_dsn,
  289             release='.'.join(map(str, VERSION)),
  290             ignore_exceptions=[
  291                 KeyboardInterrupt
  292             ]
  293         )
  294 
  295         try:
  296             run()
  297         except Exception:
  298             client.captureException()
  299 
  300     else:
  301         run()