#!/usr/bin/python
#------------------------------------------------------------------
# admin-cli-proxy-xr_run_cmd - Run admin commands from XR
#
# XR script called whenever issuing admin commands or accessing admin
#
# Angelique De Castro, June 2017
#
# Copyright (c) 2017-2018 by Cisco Systems, Inc.
# All rights reserved.
#------------------------------------------------------------------

"""
admin-cli-proxy-xr_run_cmd: A script for conforming
eXR commands issued through cXR.

This script will remove disabled console messages from admin output on XR.
For example, after removing the sysadmin connected message in config mode,
(service sysadmin-login-banner disable) this script will read its SysDB
tuple and remove the message if disabled.

Additionally, sysadmin prompts and the admin timestamp will be removed from
admin commands issued from XR. The admin output will also inherit its width and
length format from XR.

This script is available for all platforms except Scapa and Classic.
"""

# Exit program when ctrl+c is sent during import statements and logging setup
try:
    import sys
    import os
    import re
    import subprocess
    import fcntl
    import select
    import time
    import signal
    import logging
    import copy
    import termios
    import traceback

    # Logging setup
    logger = logging.getLogger("ADMIN-CLI")
    handler = logging.StreamHandler()
    formatter = logging.Formatter( \
        '%(asctime)s : %(name)s: %(levelname)s: %(message)s', \
        datefmt="%b %d %H:%M:%S")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    # Set to DEBUG during development. Set to ERROR before commit.
    logger.setLevel(logging.ERROR)

    # Suppress traceback during program execution except when log level is DEBUG
    def excepthook(excepttype, value, tb):
        tb = ''.join(traceback.format_exception(excepttype, value, tb))
        logger.warning('\n'+tb)
        sys.exit(1)
    sys.excepthook = excepthook
except KeyboardInterrupt:
    sys.exit(1)

# List of prefixes of commands that should be aborted
ABORT_CMDS = [
    "at", "ac", "conf", "deb", "in", "mo", "ru",
    "to", "upg", "term", "ti", "exit"
]

REGEX_ABORT_CMDS = [
    r"copy.*tftp.*"
]

# Flag if executing from shell without vty
NO_VTY_PRESENT = False

def restore_term(settings):
    """Restore terminal settings"""
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)

def modify_term():
    """Modify terminal settings for unbuffered stdin"""
    settings = termios.tcgetattr(sys.stdin)
    new_settings = termios.tcgetattr(sys.stdin)
    new_settings[3] = new_settings[3] & ~(termios.ECHO | termios.ICANON)
    new_settings[6][termios.VMIN] = 0
    new_settings[6][termios.VTIME] = 0
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, new_settings)
    return settings

def get_key():
    """Read a character from stdin"""
    char = os.read(sys.stdin.fileno(), 1)
    if char != None and len(char) > 0:
        return ord(char[0])

def get_term_info():
    """Get terminal dimensions"""
    # Try reading terminal length and width
    try:
        terminal = subprocess.check_output(['tty_show_line', '-t'])
        tsettings = re.search(
            r'^Length\: ([0-9]+) lines, Width\: ([0-9]+) columns$',\
            terminal, flags=re.MULTILINE)
    except Exception:
        rows, cols = 0, 0
    else:
        if tsettings:
            rows = tsettings.group(1)
            cols = tsettings.group(2)
        else:
            rows, cols = 0, 0
    return (rows, cols)

def get_suppress_info():
    """ Get which messages to suppress """
    # Perform item query
    suppress = []
    cmds = {}
    tuples = [ {'tuple':  'gl/parser/sysadmin-login-banner/disable',
                'regex':  r'1 \(0x00000001\)',
                'string': r'.* connected from [0-9]+\.[0-9]+\.[0-9]+\.[0-9]+'\
                          ' using console on .*\n'
                }]
    cmds = {'0': 'echo -e "sysdbcon -m -f /dev/stdin << EOF\n/cfg/\nget {}\n'\
                 'EOF\n" | sh | grep "\'{}\'"'.format(tuples[0]['tuple'],
                 tuples[0]['tuple'])
           }

    for index, cmd in cmds.items():
        value = None
        try:
            query = subprocess.check_output([cmd], shell=True)
            value = re.search(tuples[int(index)]['regex'], query)
        except Exception as e:
            # Set logger.setLevel(logging.DEBUG) to see this message
            logger.warning(e)
        if value is not None:
            suppress.append(tuples[int(index)]['string'])

    return suppress

def throw_error(msg):
    """Send error message to stderr and exit program"""
    logger.error(msg)
    sys.exit(1)

def check_abort(args):
    """Abort if sysadmin command should not be issued through cXR"""
    args = args.lower()
    msg = 'ABORT: Please run this command from admin prompt'
    for prefix in ABORT_CMDS:
        if args.strip().startswith(prefix):
            throw_error(msg)

    for regex in REGEX_ABORT_CMDS:
        if re.match(re.compile(regex), args):
            throw_error(msg)

def check_timestamp(content):
    """Indicate whether or not the admin timestamp is enabled"""
    if 'enable' in content:
        return 'timestamp', True

def detach_group():
    """Detach from parent process group to prevent signal propagation"""
    # Not for shell drop; possible side effects
    if sys.argv[1:]:
        os.setpgrp()

def sanitize(line, args):
    """Sanitize output of all unwanted messages"""
    suppress = copy.copy(args['suppress'])
    for msg in suppress:
        line = re.sub(msg, '', line, re.MULTILINE)
        args['suppress'].remove(msg)
    return line, args

"""Reengineered subprocess class for handling interactive commands"""

class IPopen(subprocess.Popen):
    """Subprocess steer and transreceive module"""
    SYS_PROMPT = 'sysadmin-vm:.*:?.*# $'
    INT_PROMPT = r'\??.*\[[0-9a-zA-Z/,.]+\]:?\s?(\?|\:)?\s?$'
    MOR_PROMPT = '--More--'
    POLL_INTERVAL = 0.1
    AUTH_BYP_MSG = 'Authentication is bypassed, can not connect to admin plane'

    # Page navigation character input: [space, return, 'n', 'q']
    NAV_CHARS = [32, 10, 110, 113]

    params = {
        'timestamp':        False,
        'supplied_command': None,
        'suppress':          []
    }

    def __init__(self, *args, **kwargs):
        """Initialize communication with child process"""
        command = kwargs.get('command', '')

        # Throw syntax error if command is improperly executed
        if ';' in command:
            throw_error("Invalid input detected")

        # Abort if command shouldn't be issued through cXR
        check_abort(command)

        # Keep track of supplied command
        self.params['supplied_command'] = command

        # Get messages to suppress
        self.params['suppress'] = get_suppress_info()

        kwargs_ = {
            # stdin not PIPE'ed for shell drops
            'stdin': subprocess.PIPE if command else None,
            'stdout': subprocess.PIPE,
            'stderr': subprocess.STDOUT,
            'preexec_fn': detach_group,
            'close_fds': True,
            'universal_newlines': True,
            'bufsize': 0
        }

        subprocess.Popen.__init__(self, *args, **kwargs_)

        # Make all output fds non-blocking
        for fd_o in (self.stdout, self.stderr):
            if fd_o:
                fcntl.fcntl(
                    fd_o.fileno(),
                    fcntl.F_SETFL,
                    fcntl.fcntl(fd_o.fileno(), fcntl.F_GETFL) | os.O_NONBLOCK
                )

        # Wait for the initial prompt
        self.await([self.SYS_PROMPT], enable=command, display=(not command))

        # Code below does not apply for shell drop
        if command:
            # Honor XR's terminal settings. Prepare command here.
            if NO_VTY_PRESENT == False:
                term_config = \
                    "terminal length {} ; terminal width {}"\
                    .format(*get_term_info())

                # Pass on XR's terminal settings
                # Check if sysadmin timestamp enabled
                self.interact(term_config + " ; show cli | include timestamp", \
                            close=False, display=False, callback=check_timestamp)
            else:
                self.interact(" show cli | include timestamp", \
                          close=False, display=False, callback=check_timestamp)


            # Register signal handler before interaction
            signal.signal(signal.SIGINT, self.signal_handler)

            # Determine if this is a help request
            terminate = not command.endswith('?')

            # Begin the interaction
            self.interact(command, terminate=terminate)

        # Wait for the process to end
        self.wait()

    def signal_handler(self, sig, frame):
        """Hand-propagation of SIGINT to child process"""
        # Send the ascii code for ctrl-c as \003
        self.interact('\003', close=True, terminate=False)

    def await(self, prompts, enable=True, display=True, command='', lead=None,\
              callback=None, terminate=False):
        """Receive stdout, sanitize and wait for prompts"""
        cmd_removed = False
        term_removed = False
        prompt_found = False
        time_removed = False

        return_prompt = (self.SYS_PROMPT)[:-1] + re.escape(lead) + \
                        '[ \r\n$]*' if lead else self.SYS_PROMPT

        params = self.params

        if lead:
            prompts[prompts.index(self.SYS_PROMPT)] = return_prompt

        try:
            # Execute commands unless child process has terminated
            while self.poll() is None:
                # Monitor stdout for read content every 0.1 seconds
                if select.select([self.stdout.fileno()], [], [], 0.1)[0]:
                    chunk = self.stdout.read()

                    # Exit if user is not a legitimate sysadmin user
                    if self.AUTH_BYP_MSG in chunk:
                        throw_error(self.AUTH_BYP_MSG)

                    if callback:
                        callback_tuple = callback(chunk)
                        if callback_tuple:
                            params[callback_tuple[0]] = callback_tuple[1]

                    # Perform default noise removal
                    if display:
                        if enable:
                            if command:
                                if not cmd_removed:
                                    # Remove command from output
                                    command_ = command[:-1] if lead else command
                                    if chunk.startswith(command_):
                                        chunk = chunk[len(command_):]
                                        cmd_removed = True
                                if cmd_removed and not term_removed:
                                    # Remove terminator from output
                                    if chunk and ((lead and chunk[0] == '?') or \
                                    (not lead and chunk[0] == '\n')):
                                        chunk = chunk[1:]
                                        term_removed = True
                                if term_removed and not time_removed and \
                                params['timestamp']:
                                    # Remove the timestamp and its terminator
                                    chunk, nsubs = re.subn(
                                            r"^(?P<tsstr>(?P<datestr>(\w{3})[ ]+"\
                                            "(\w{3})[ ]+(\d{1,2})[ ]+(\d{2}):(\d"\
                                            "{2}):(\d{2}))(.(?P<milli>\d{3}))?( "\
                                            "(?P<tz>\w+))?)(?=\r|\n|$|\r\n)",\
                                            "", chunk, count=1, flags=re.MULTILINE)
                                    if nsubs == 1:
                                        if chunk.startswith('\n'):
                                            chunk = chunk[1:]
                                        time_removed = True

                                chunk, params = sanitize(chunk, params)

                            # Remove the return prompt
                            if re.search(return_prompt, chunk, re.MULTILINE):
                                chunk = re.sub(return_prompt, '', chunk, \
                                                re.MULTILINE)
                                prompt_found = True
                                chunk, params = sanitize(chunk, params)
                                chunk = re.sub(return_prompt, '', chunk, \
                                               re.MULTILINE)

                        # Sanitize output of admin commands
                        if not enable and not term_removed:
                            chunk, params = sanitize(chunk, params)

                        sys.stdout.write(chunk)
                        sys.stdout.flush()

                    # Catch a prompt and yield
                    if enable:
                        for index, prompt in enumerate(prompts):
                            if (prompt_found and prompt == return_prompt) or \
                            re.search(prompt, chunk):
                                return index
                # Skip return branch if entering admin mode or executing
                # non-commad help admin commands
                elif (not command) or terminate:
                    continue
                # If stdout is done reading but sysadmin prompt hasn't 
                # returned, exit from child process
                elif select.select([], [self.stdin.fileno()], [])[1] and command:
                    return 1
        except KeyboardInterrupt:
            # Graceful exit
            sys.exit(1)

    def interact(self, command, stdin='', close=True, terminate=True, \
                 display=True, callback=None):
        """Communicate and handle prompts"""
        # Check if child process is still running
        if self.poll() is not None:
            sys.exit(1)

        # Stdin could be the supplied command or a response to a prompt
        if not stdin:
            stdin = command
        # Remove the prompt and stdin echo from stdout
        remove = stdin
        # Supplied command does not end with ?, so add a newline to the command
        if terminate:
            stdin = stdin + '\n'

        params = self.params

        # Kill child process and return the prompt
        if stdin == '\003':
            command = params['supplied_command']

        # Initiate communication
        self.stdin.write(stdin)
        self.stdin.flush()

        # Execute only while child process is still running
        while self.poll() is None:
            # Hunt for prompts
            p_index = self.await(
                [self.INT_PROMPT, self.SYS_PROMPT, self.MOR_PROMPT],
                enable=True,
                display=display,
                command=remove,
                lead=command[:-1] if command.endswith('?') else None,
                callback=callback,
                terminate=terminate
            )
            if p_index == 0:
                if '?' in params['supplied_command']:
                    self.interact('\003', close=True, terminate=False)
                # Child process needs user input
                try:
                    user_input = raw_input()
                # Send sigint to admin if user hits ctrl-d
                except:
                    self.signal_handler(None, None)
                self.interact(params['supplied_command'], stdin=user_input)
                return
            elif p_index == 1:
                # sysadmin prompt received; close & clean up
                if close:
                    sys.exit()
                return
            elif p_index == 2:
                # Page navigation prompt received
                settings = modify_term()
                while True:
                    user_input = get_key()
                    if user_input in self.NAV_CHARS:
                        sys.stdout.write('\r')
                        sys.stdout.flush()
                        self.interact(params['supplied_command'], \
                                      stdin=chr(user_input), \
                                      close=True, terminate=False)
                        restore_term(settings)
                        return
                    else:
                        time.sleep(self.POLL_INTERVAL)


if __name__ == "__main__":
    # External admin shell proxy executable
    admin_proxy = ['/pkg/sbin/admin-cli-proxy-xr_static', '-p', 'XR']

    # Seed command received from user 
    cmd = ' '.join(sys.argv[1:])
    if sys.argv[-1] == "-no_vty_present":
        NO_VTY_PRESENT = True
        cmd = " ".join(sys.argv[1:-1])

    IPopen(admin_proxy, command=cmd)
