<commit_before>
__author__    = "<NAME>"
__copyright__ = "Copyright 2012-2013, The SAGA Project"
__license__   = "MIT"


import re
import os
import sys
import errno

import saga.utils.logger
import saga.utils.which
import saga.utils.pty_shell_factory as supsf

_PTY_TIMEOUT = 2.0
_SCHEMAS     = ['ssh', 'gsissh', 'fork', 'shell', 'file', 'scp', 'sftp', 'gsiscp', 'gsisftp']
# FIXME: gsiftp?

# ------------------------------------------------------------------------------
#
# iomode flags
#
IGNORE   = 0    # discard stdout / stderr
MERGED   = 1    # merge stdout and stderr
SEPARATE = 2    # fetch stdout and stderr individually (one more hop)
STDOUT   = 3    # fetch stdout only, discard stderr
STDERR   = 4    # fetch stderr only, discard stdout


# --------------------------------------------------------------------
#
class PTYShell (object) :
    """
    This class wraps a shell process and runs it as a :class:`PTYProcess`.  The
    user of this class can start that shell, and run arbitrary commands on it.

    The shell to be run is expected to be POSIX compliant (bash, dash, sh, ksh
    etc.) -- in particular, we expect the following features:
    ``$?``,
    ``$!``,
    ``$#``,
    ``$*``,
    ``$@``,
    ``$$``,
    ``$PPID``,
    ``>&``,
    ``>>``,
    ``>``,
    ``<``,
    ``|``,
    ``||``,
    ``()``,
    ``&``,
    ``&&``,
    ``wait``,
    ``kill``,
    ``nohup``,
    ``shift``,
    ``export``,
    ``PS1``, and
    ``PS2``.

    Note that ``PTYShell`` will change the shell prompts (``PS1`` and ``PS2``),
    to simplify output parsing.  ``PS2`` will be empty, ``PS1`` will be set
    ``PROMPT-$?->`` -- that way, the prompt will report the exit value of the
    last command, saving an extra roundtrip.  Users of this class should be
    careful when setting other prompts -- see :func:`set_prompt` for more
    details.

    Usage Example::
    ^^^^^^^^^^^^^^^

        # start the shell, find its prompt.  
        self.shell = saga.utils.pty_shell.PTYShell ("ssh://user@remote.host.net/", contexts, self._logger)

        # run a simple shell command, merge stderr with stdout.  $$ is the pid
        # of the shell instance.
        ret, out, _ = self.shell.run_sync ("mkdir -p /tmp/data.$$/" )

        # check if mkdir reported success
        if  ret != 0 :
            raise saga.NoSuccess ("failed to prepare base dir (%s)(%s)" % (ret, out))

        # stage some data from a local string variable into a file on the remote system
        self.shell.stage_to_remote (src = pbs_job_script, 
                                    tgt = "/tmp/data.$$/job_1.pbs")

        # check size of staged script (this is actually done on PTYShell level
        # already, with no extra hop):
        ret, out, _ = self.shell.run_sync ("stat -c '%s' /tmp/data.$$/job_1.pbs" )
        if  ret != 0 :
            raise saga.NoSuccess ("failed to check size (%s)(%s)" % (ret, out))

        assert (len(pbs_job_script) == int(out))


    Data Staging and Data Management:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    The PTYShell class does not only support command execution, but also basic
    data management: for SSH based shells, it will create a tunneled scp/sftp
    connection for file staging.  Other data management operations (mkdir, size,
    list, ...) are executed either as shell commands, or on the scp/sftp channel
    (if possible on the data channel, to keep the shell pty free for concurrent
    command execution).  Ssh tunneling is implemented via ssh.v2 'ControlMaster'
    capabilities (see `ssh_config(5)`).
    
    For local shells, PTYShell will create an additional shell pty for data
    management operations.  


    Asynchronous Notifications:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^

    A third pty process will be created for asynchronous notifications.  For
    that purpose, the shell started on the first channel will create a named
    pipe, at::

      $HOME/.saga/adaptors/shell/async.$$

    ``$$`` here represents the pid of the shell process.  It will also set the
    environment variable ``SAGA_ASYNC_PIPE`` to point to that named pipe -- any
    application running on the remote host can write event messages to that
    pipe, which will be available on the local end (see below).  `PTYShell`
    leaves it unspecified what format those messages have, but messages are
    expected to be separated by newlines.
    
    An adaptor using `PTYShell` can subscribe for messages via::

      self.pty_shell.subscribe (callback)

    where callback is a Python callable.  PTYShell will listen on the event
    channel *in a separate thread* and invoke that callback on any received
    message, passing the message text (sans newline) to the callback.

    An example usage: the command channel may run the following command line::

      ( sh -c 'sleep 100 && echo "job $$ done" > $SAGA_ASYNC_PIPE" \
                         || echo "job $$ fail" > $SAGA_ASYNC_PIPE" ) &

    which will return immediately, and send a notification message at job
    completion.

    Note that writes to named pipes are not atomic.  From POSIX:

    ``A write is atomic if the whole amount written in one operation is not
    interleaved with data from any other process. This is useful when there are
    multiple writers sending data to a single reader. Applications need to know
    how large a write request can be expected to be performed atomically. This
    maximum is called {PIPE_BUF}. This volume of IEEE Std 1003.1-2001 does not
    say whether write requests for more than {PIPE_BUF} bytes are atomic, but
    requires that writes of {PIPE_BUF} or fewer bytes shall be atomic.`

    Thus the user is responsible for ensuring that either messages are smaller
    than *PIPE_BUF* bytes on the remote system (usually at least 1024, on Linux
    usually 4096), or to lock the pipe on larger writes.


    Automated Restart, Timeouts:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    For timeout and restart semantics, please see the documentation to the
    underlying :class:`saga.utils.pty_process.PTYProcess` class.

    """

    # TODO: 
    #   - on client shell activitites, also mark the master as active, to
    #     avoid timeout garbage collection.
    #   - use ssh mechanisms for master timeout (and persist), as custom
    #     mechanisms will interfere with gc_timout.

    # ----------------------------------------------------------------
    #
    def __init__ (self, url, session, logger=None, init=None, opts={}) :

        self.url       = url      # describes the shell to run
        self.logger    = logger   # possibly log to here
        self.init      = init     # call after reconnect
        self.opts      = opts     # options...

        self.prompt          = None
        self.prompt_re       = None
        self.initialize_hook = None
        self.finalize_hook   = None

        # we need a local dir for file staging caches.  At this point we use
        # $HOME, but should make this configurable (FIXME)
        self.base = os.environ['HOME'] + '/.saga/adaptors/shell/'

        try:
            os.makedirs (self.base)

        except OSError as e:
            if e.errno == errno.EEXIST and os.path.isdir (self.base):
                pass
            else: 
                raise saga.NoSuccess ("could not create staging dir: %s" % e)

        
        # need a new logger?
        if not self.logger :
            self.logger = saga.utils.logger.getLogger ('PTYShell')

        self.factory    = supsf.PTYShellFactory   ()
        self.pty_info   = self.factory.initialize (url, session, self.logger)
        self.pty_shell  = self.factory.run_shell  (self.pty_info)

        self.pty_shell.set_initialize_hook (self.initialize)
        self.pty_shell.set_finalize_hook   (self.finalize)

        self.initialize ()


    # ----------------------------------------------------------------
    #
    def __del__ (self) :

        self.finalize (kill_pty=True)


    # ----------------------------------------------------------------------
    #
    def set_initialize_hook (self, initialize_hook) :

        self.initialize_hook = initialize_hook


    # ----------------------------------------------------------------------
    #
    def set_finalize_hook (self, finalize_hook) :

        self.finalize_hook = finalize_hook


    # ----------------------------------------------------------------
    #
    def initialize (self) :
        """ initialize the shell connection.  """

        # FIXME: timout should be something like n*latency

        # run a POSIX compatible shell, usually /bin/sh, in interactive mode
        # also, turn off tty echo
        command_shell = "exec /bin/sh -i"

        # use custom shell if so requested
        if  'shell' in self.opts and self.opts['shell'] :
            command_shell = "exec %s" % self.opts['shell']
            self.logger.info ("custom  command shell: %s" % command_shell)


        self.logger.debug    ("running command shell: %s" % command_shell)
        self.pty_shell.write ("stty -echo ; %s\n"         % command_shell)

        # make sure this worked, and that we find the prompt. We use
        # a versatile prompt pattern to account for the custom shell case.

        # find initial prompt from login shell
        ret1, out1 = self.pty_shell.find (["^.*\$\s*$"], 10.0)

        # we need to make sure that the found prompt is really the prompt of
        # the new shell, so we check $? for successfull startup.
        self.pty_shell.write ("echo $?\n")
        ret2, out2 = self.pty_shell.find (["\d+\n"],     10.0)
        ret3, out3 = self.pty_shell.find (["^.*\$\s*$"], 10.0)

        if  int(out2) != 0 :
            raise saga.NoSuccess ("cannot initalize shell with %s (%s)" \
                    % (self.opts['shell'], out1))


        # set and register new prompt
        self.run_sync ( "unset PROMPT_COMMAND ; "
                             + "PS1='PROMPT-$?->'; "
                             + "PS2=''; "
                             + "export PS1 PS2\n", <commit_msg>fix /bin/sh startup check on non-default sh prompt<commit_after>
__author__    = "<NAME>"
__copyright__ = "Copyright 2012-2013, The SAGA Project"
__license__   = "MIT"


import re
import os
import sys
import errno

import saga.utils.logger
import saga.utils.which
import saga.utils.pty_shell_factory as supsf

_PTY_TIMEOUT = 2.0
_SCHEMAS     = ['ssh', 'gsissh', 'fork', 'shell', 'file', 'scp', 'sftp', 'gsiscp', 'gsisftp']
# FIXME: gsiftp?

# ------------------------------------------------------------------------------
#
# iomode flags
#
IGNORE   = 0    # discard stdout / stderr
MERGED   = 1    # merge stdout and stderr
SEPARATE = 2    # fetch stdout and stderr individually (one more hop)
STDOUT   = 3    # fetch stdout only, discard stderr
STDERR   = 4    # fetch stderr only, discard stdout


# --------------------------------------------------------------------
#
class PTYShell (object) :
    """
    This class wraps a shell process and runs it as a :class:`PTYProcess`.  The
    user of this class can start that shell, and run arbitrary commands on it.

    The shell to be run is expected to be POSIX compliant (bash, dash, sh, ksh
    etc.) -- in particular, we expect the following features:
    ``$?``,
    ``$!``,
    ``$#``,
    ``$*``,
    ``$@``,
    ``$$``,
    ``$PPID``,
    ``>&``,
    ``>>``,
    ``>``,
    ``<``,
    ``|``,
    ``||``,
    ``()``,
    ``&``,
    ``&&``,
    ``wait``,
    ``kill``,
    ``nohup``,
    ``shift``,
    ``export``,
    ``PS1``, and
    ``PS2``.

    Note that ``PTYShell`` will change the shell prompts (``PS1`` and ``PS2``),
    to simplify output parsing.  ``PS2`` will be empty, ``PS1`` will be set
    ``PROMPT-$?->`` -- that way, the prompt will report the exit value of the
    last command, saving an extra roundtrip.  Users of this class should be
    careful when setting other prompts -- see :func:`set_prompt` for more
    details.

    Usage Example::
    ^^^^^^^^^^^^^^^

        # start the shell, find its prompt.  
        self.shell = saga.utils.pty_shell.PTYShell ("ssh://user@remote.host.net/", contexts, self._logger)

        # run a simple shell command, merge stderr with stdout.  $$ is the pid
        # of the shell instance.
        ret, out, _ = self.shell.run_sync ("mkdir -p /tmp/data.$$/" )

        # check if mkdir reported success
        if  ret != 0 :
            raise saga.NoSuccess ("failed to prepare base dir (%s)(%s)" % (ret, out))

        # stage some data from a local string variable into a file on the remote system
        self.shell.stage_to_remote (src = pbs_job_script, 
                                    tgt = "/tmp/data.$$/job_1.pbs")

        # check size of staged script (this is actually done on PTYShell level
        # already, with no extra hop):
        ret, out, _ = self.shell.run_sync ("stat -c '%s' /tmp/data.$$/job_1.pbs" )
        if  ret != 0 :
            raise saga.NoSuccess ("failed to check size (%s)(%s)" % (ret, out))

        assert (len(pbs_job_script) == int(out))


    Data Staging and Data Management:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    The PTYShell class does not only support command execution, but also basic
    data management: for SSH based shells, it will create a tunneled scp/sftp
    connection for file staging.  Other data management operations (mkdir, size,
    list, ...) are executed either as shell commands, or on the scp/sftp channel
    (if possible on the data channel, to keep the shell pty free for concurrent
    command execution).  Ssh tunneling is implemented via ssh.v2 'ControlMaster'
    capabilities (see `ssh_config(5)`).
    
    For local shells, PTYShell will create an additional shell pty for data
    management operations.  


    Asynchronous Notifications:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^

    A third pty process will be created for asynchronous notifications.  For
    that purpose, the shell started on the first channel will create a named
    pipe, at::

      $HOME/.saga/adaptors/shell/async.$$

    ``$$`` here represents the pid of the shell process.  It will also set the
    environment variable ``SAGA_ASYNC_PIPE`` to point to that named pipe -- any
    application running on the remote host can write event messages to that
    pipe, which will be available on the local end (see below).  `PTYShell`
    leaves it unspecified what format those messages have, but messages are
    expected to be separated by newlines.
    
    An adaptor using `PTYShell` can subscribe for messages via::

      self.pty_shell.subscribe (callback)

    where callback is a Python callable.  PTYShell will listen on the event
    channel *in a separate thread* and invoke that callback on any received
    message, passing the message text (sans newline) to the callback.

    An example usage: the command channel may run the following command line::

      ( sh -c 'sleep 100 && echo "job $$ done" > $SAGA_ASYNC_PIPE" \
                         || echo "job $$ fail" > $SAGA_ASYNC_PIPE" ) &

    which will return immediately, and send a notification message at job
    completion.

    Note that writes to named pipes are not atomic.  From POSIX:

    ``A write is atomic if the whole amount written in one operation is not
    interleaved with data from any other process. This is useful when there are
    multiple writers sending data to a single reader. Applications need to know
    how large a write request can be expected to be performed atomically. This
    maximum is called {PIPE_BUF}. This volume of IEEE Std 1003.1-2001 does not
    say whether write requests for more than {PIPE_BUF} bytes are atomic, but
    requires that writes of {PIPE_BUF} or fewer bytes shall be atomic.`

    Thus the user is responsible for ensuring that either messages are smaller
    than *PIPE_BUF* bytes on the remote system (usually at least 1024, on Linux
    usually 4096), or to lock the pipe on larger writes.


    Automated Restart, Timeouts:
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^

    For timeout and restart semantics, please see the documentation to the
    underlying :class:`saga.utils.pty_process.PTYProcess` class.

    """

    # TODO: 
    #   - on client shell activitites, also mark the master as active, to
    #     avoid timeout garbage collection.
    #   - use ssh mechanisms for master timeout (and persist), as custom
    #     mechanisms will interfere with gc_timout.

    # ----------------------------------------------------------------
    #
    def __init__ (self, url, session, logger=None, init=None, opts={}) :

        self.url       = url      # describes the shell to run
        self.logger    = logger   # possibly log to here
        self.init      = init     # call after reconnect
        self.opts      = opts     # options...

        self.prompt          = None
        self.prompt_re       = None
        self.initialize_hook = None
        self.finalize_hook   = None

        # we need a local dir for file staging caches.  At this point we use
        # $HOME, but should make this configurable (FIXME)
        self.base = os.environ['HOME'] + '/.saga/adaptors/shell/'

        try:
            os.makedirs (self.base)

        except OSError as e:
            if e.errno == errno.EEXIST and os.path.isdir (self.base):
                pass
            else: 
                raise saga.NoSuccess ("could not create staging dir: %s" % e)

        
        # need a new logger?
        if not self.logger :
            self.logger = saga.utils.logger.getLogger ('PTYShell')

        self.factory    = supsf.PTYShellFactory   ()
        self.pty_info   = self.factory.initialize (url, session, self.logger)
        self.pty_shell  = self.factory.run_shell  (self.pty_info)

        self.pty_shell.set_initialize_hook (self.initialize)
        self.pty_shell.set_finalize_hook   (self.finalize)

        self.initialize ()


    # ----------------------------------------------------------------
    #
    def __del__ (self) :

        self.finalize (kill_pty=True)


    # ----------------------------------------------------------------------
    #
    def set_initialize_hook (self, initialize_hook) :

        self.initialize_hook = initialize_hook


    # ----------------------------------------------------------------------
    #
    def set_finalize_hook (self, finalize_hook) :

        self.finalize_hook = finalize_hook


    # ----------------------------------------------------------------
    #
    def initialize (self) :
        """ initialize the shell connection.  """

        # FIXME: timout should be something like n*latency

        # run a POSIX compatible shell, usually /bin/sh, in interactive mode
        # also, turn off tty echo
        command_shell = "exec /bin/sh -i"

        # use custom shell if so requested
        if  'shell' in self.opts and self.opts['shell'] :
            command_shell = "exec %s" % self.opts['shell']
            self.logger.info ("custom  command shell: %s" % command_shell)


        self.logger.debug    ("running command shell: %s" % command_shell)
        self.pty_shell.write ("stty -echo ; %s\n"         % command_shell)

        # make sure this worked, and that we find the prompt. We use
        # a versatile prompt pattern to account for the custom shell case.

        # find initial prompt from login shell
        ret1, out1 = self.pty_shell.find (["^.*\$\s*$"], 10.0)

        # we need to make sure that the found prompt is really the prompt of
        # the new shell, so we check $? for successfull startup.
        self.pty_shell.write ('echo "-SAGA-$?"\n')
        ret2, out2 = self.pty_shell.find (["-SAGA-"],    10.0)
        ret3, out3 = self.pty_shell.find (["\d+\n"],     10.0)
        ret4, out4 = self.pty_shell.find (["^.*\$\s*$"], 10.0)

        if  int(out3[:-1]) != 0 :
            raise saga.NoSuccess ("cannot initalize shell with %s (%s)" \
                    % (self.opts['shell'], out1))


        # set and register new prompt
        self.run_sync ( "unset PROMPT_COMMAND ; "
                             + "PS1='PROMPT-$?->'; "
                             + "PS2=''; "
                             + "export PS1 PS2\n", 
{'fraction_non_alphanumeric': 0.15334687608383293, 'fraction_numerical': 0.0071347173363721945, 'mean_word_length': 2.0035714285714286, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3184180', 'n_tokens_mistral': 5960, 'n_tokens_neox': 5622, 'n_words': 2102}