Python Bash interaction -- ALPHA (at best)

python_bash.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Interact with bash from python

ALPHA at best and hasn't been tested in a while
"""
from __future__ import print_function, unicode_literals

import subprocess
import random
import time
from select import select

class TimedoutError(Exception):
    pass

class Bash(object):
    """
    Start a bash processes that Python can interact with

    Options: 
    --------
    (These, except for err2out, can also be set directly and/or changed)

    timeout : [None]
        Specify how long to wait for input before timing out and returning.
    
        The timeout is between lines and *not* the entire processes
    
        Note that if a timeout occurs, the error stream will end with:
            >>STDOUT_TIMEOUT<<  --OR-- >>STDERR_TIMEOUT<<
    
        A None will not have a timeout

    err2out : [False]
        If True, will stream stderr to stdout. Otherwise, it will
        be its own separate stream. Note that calling read(err=True) or 
        communicate(err=True) with err2out=True will just return empty
        stderrs and will also never call stderr_fcn.

    stdout_fcn / stderr_fcn: [None]
        Specify a callable object which will get called on every line
        of stdout/stderr as it is streamed

    Example:
    -------
    This is not designed to be soley an SSH tool, but one of its strengths is 
    that is can opperate as such
    
        # For this example, we want to print to screen to see the output
        def print_(txt):
            print('> ' + txt)
        
        proc = Bash(timeout=10,err2out=True,stdout_fcn=print_)
        
    Notes and Issues:
    -----------------
    * This will read from stdout first and then after that closes 
      (or times out), it will read stderr
    * Once a timeout occurs, it will raise a TimedoutError upon future
      calls. To override this, set `proc._timedout=False` but BE WARNED
      that the next call will likely caputure remaining output

    WARNING:
    --------
    This is *not* intended for user-specified code or input. There is zero
    parsing or protection from malicious input. 

    """
    def __init__(self,timeout=None,err2out=False,\
                 stdout_fcn=None,stderr_fcn=None):
        """
        Initialize the bash subprocess
        """
        
        stderr = subprocess.PIPE
        self._err2out = err2out
        if err2out:
            stderr = subprocess.STDOUT
        
        self.proc = subprocess.Popen(['/usr/bin/env','bash'], stdin=subprocess.PIPE, 
                    stdout=subprocess.PIPE,stderr=stderr,shell=False)
    
        self.timeout = timeout
        self.stdout_fcn = stdout_fcn
        self.stderr_fcn = stderr_fcn
        self._timedout = False
    
    def readlines(self):
        """
        Main function to read the stdout and stderr stream.
        
        Works by sending a randomly-generated sentinel text to denote
        the end of the stream
        
        Returns
            stdout,stderr -- Lists with output. 
        
        Note: Use `read` to read as a single block of text
        """
        if self._timedout:
            raise TimedoutError('Prev. command timed out. Set _timedout=False to override') 
    
        end_txt = ''.join(random.choice('abcdefghijklmnopqrstuvwxyz') \
                    for _ in xrange(30))
        self.end_txt = end_txt
    
        self.write("echo "+end_txt) # send end_txt
        result = []
        line,_ = self._stdout()
        while line.strip()!=end_txt: # look for end_txt
            result.append(line)
            line,ready = self._stdout()
            if not ready:
                self._timedout = True
                if self._err2out:
                    return result + ['>>STDOUT_TIMEOUT<<'],[]
                else:
                    return result,['>>STDOUT_TIMEOUT<<']
    
        if self._err2out:
            return result,[]     
       
        self.write(">&2 echo "+end_txt) # send end_txt
        result_err = []
        line,_ = self._stderr()
        while line.strip()!=end_txt: # look for end_txt
            result_err.append(line)
            line,ready = self._stderr()
            if not ready:
                self._timedout = True
                return result,result_err + ['>>STDOUT_TIMEOUT<<']
            
        return result,result_err

    def write(self, data):
        """
        Send `data` to the underlying bash subprocess
        """
        if not data.endswith('\n'):
            data += '\n'
        self.proc.stdin.write(data)
    def read(self,err=False):
        """
        Read the stream.
        
        If err=False, returns stdout string
        if err=True, returns (stdout,stderr) strings
        
        Recall that if stdout_fcn and stderr_fcn will be called upon read
        """
        result,result_err = self.readlines()
        if err:
            return '\n'.join(result),'\n'.join(result_err)
        return '\n'.join(result)

    def communicate(self, data,err=False,sleep=None):
        """
        Combined write and read in one function call.
        
        Setting `sleep` is useful to help buffer the commincation
        """
        self.write(data)
        if sleep is not None:
            time.sleep(sleep)
        return self.read(err=err)
    
    def kill(self):
        """
        Close (kill) the bash process
        """
        self.proc.kill()

    def _stdout(self):
        _ready, _, _ = select([self.proc.stdout], [], [], self.timeout)
        if _ready:
            line = self.proc.stdout.readline().rstrip()
            ready = True
            if hasattr(self.stdout_fcn,'__call__') and line.strip() !=self.end_txt:
                self.stdout_fcn(line)
        else:
            line = ''
            ready = False
        return line,ready

    def _stderr(self):
        _ready, _, _ = select([self.proc.stderr], [], [], self.timeout)
        if _ready:
            line = self.proc.stderr.readline().rstrip()
            ready = True
            if hasattr(self.stderr_fcn,'__call__') and line.strip() !=self.end_txt:
                self.stderr_fcn(line)
        else:
            line = ''
            ready = False
        return line,ready