connection.py

import socket
import select
import subprocess
import re

import logging
l = logging.getLogger('puppeteer.connection')
#l.setLevel(logging.DEBUG)

from .errors import ConnectionFail

def rw_alias(f):
    '''
    Makes an alias for read/recv and write/send.
    '''
    print "CLASS:",f.im_class
    return f

class Connection(object):
    '''
    A connection handler for puppeteer. Basically a wrapper around sockets or
    files.

    Will handle all sorts of intelligent stuff like timeouts and crap.
    '''
    def __init__(self, host=None, port=None, exe=None, args=(), s=None, fd=None):
        self.host = host
        self.port = port
        self.exe_args = [ exe ] + list(args)

        self.s = s
        self.p = fd
        self.connected = None if s is None and fd is None else True

    def copy(self):
        return Connection(host=self.host, port=self.port, exe=self.exe_args[0], args=self.exe_args[1:])

    def connect(self):
        '''
        Connect!
        '''
        if self.host is not None and self.port is not None:
            self.s = socket.create_connection((self.host, self.port))
        elif self.exe_args[0] is not None:
            self.p = subprocess.Popen(self.exe_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=0)
        else:
            raise Exception("How are we supposed to connect, man?")

        l.info("Connected!")
        self.connected = True
        return self

    def close(self):
        if self.s:
            self.s.close()
        elif self.p:
            self.p.close()

    def send(self, msg):
        '''
        Send the message.
        '''
        l.debug("Sending: %r", msg)

        try:
            if self.s is not None:
                return self.s.sendall(msg)
            elif self.p is not None:
                return self.p.stdin.write(msg)
        except socket.error as e:
            raise ConnectionFail(str(e))

    def recv(self, n, timeout=None, quash=False):
        '''
        Receive up to n bytes.
        '''
        #l.debug("recv(%d, timeout=%s)", n, timeout)
        slist = [ self.s if self.s is not None else self.p.stdout ]
        if timeout is not None:
            (rlist, _, _) = select.select(slist, [], [], timeout)
        else:
            (rlist, _, _) = select.select(slist, [], [])

        l.debug("recv with size %d and timeout %s", n, timeout)

        #print rlist, wlist, xlist
        if rlist == []:
            l.debug("RECV TIMEOUT")
            if not quash:
                raise ConnectionFail("very timeout")
            else:
                return ""
        try:
            readsock = rlist[0]
            if type(readsock) == socket.socket:
                r = readsock.recv(n)
            if type(readsock) == file:
                r = readsock.read(n)

            l.debug("read: %r",r)
        except socket.error:
            r = ""

        if len(r) == 0 and n != 0:
            raise ConnectionFail("Received nothing. Much sad.")

        #l.debug("Got: %r", r)
        return r

    def read(self, n=None, timeout=None, quash=False):
        '''
        Read exactly n bytes, or all that's available if n is None.
        '''
        if n is None:
            self.read_all(timeout=timeout)

        result = ""
        while len(result) < n:
            tmp = self.recv(n - len(result), timeout=timeout, quash=quash)
            if tmp == "" and timeout is not None:
                break
            result += tmp

        return result

    # read until the given string
    def read_until(self, c=None, regex=None, max_chars=None, timeout=None):
        '''
        Read until the given string, or matching the given regex.
        '''
        if c is None and regex is None:
            # they cannot both be None
            raise Exception('You must specify at least one argument')
        if c is not None and regex is not None:
            raise Exception('You cannot specify both c and regex.')

        if c is not None:
            l.debug("Reading until: %r", c)
        else:
            l.debug("Reading until: %s", regex)

        buf = ""
        tmp = None
        while max_chars is None or len(buf) < max_chars:
            l.debug("... so far: %s", buf)
            if c is not None and buf.endswith(c):
                break
            elif regex is not None and re.search(regex, buf):
                break

            tmp = self.recv(1, timeout=timeout)
            if tmp == "":
                break
            buf += tmp
        l.debug("... read: %r", buf)
        return buf

    def read_all(self, timeout=None, stepsize=8192):
        '''
        Read as long as there is more stuff and the timeout is not expired.
        '''
        buff = ""

        while True:
            try:
                s = self.recv(stepsize, timeout=timeout)
            except ConnectionFail:
                break

            if len(s) == 0:
                break
            buff += s

        return buff

    def shutdown(self, what):
        if self.s is not None:
            self.s.shutdown(what)
        elif self.p is not None:
            if what == socket.SHUT_WR:
                self.p.stdin.close()
            if what == socket.SHUT_RD:
                self.p.stdout.close()