tcp.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import socket
import time
import sys

from Peach.Engine.engine import Engine
from Peach.Engine.common import PeachException
from Peach.publisher import Publisher
from Peach.publisher import Timeout
from Peach.publisher import PublisherSoftException
from Peach.Utilities.common import *
import Peach


def Debug(msg):
    if Peach.Engine.engine.Engine.debug:
        print(msg)

#class Timeout(SoftException):
#	def __init__(self, msg):
#		self.msg = msg
#	
#	def __str__(self):
#		return self.msg

class Tcp(Publisher):
    """
    A simple TCP client publisher.
    """


    def __init__(self, host, port, timeout=0.25, throttle=0):
        """
        @type	host: string
        @param	host: Remote host
        @type	port: number
        @param	port: Remote port
        @type	timeout: number
        @param	timeout: How long to wait for reponse
        @type	throttle: number
        @param	throttle: How long to wait between connections
        """
        Publisher.__init__(self)
        self._host = host

        try:
            self._port = int(port)
        except:
            raise PeachException("The Tcp publisher parameter for port was not a valid number.")

        try:
            self._timeout = float(timeout)
        except:
            raise PeachException("The Tcp publisher parameter for timeout was not a valid number.")

        try:
            self._throttle = float(throttle)
        except:
            raise PeachException("The Tcp publisher parameter for throttle was not a valid number.")

        self._socket = None

    def start(self):
        pass

    def stop(self):
        self.close()

    def connect(self):
        """
        Create connection.
        """
        self.close()

        if self._throttle > 0:
            time.sleep(self._throttle)

        # Try connecting many times
        # before we crash.
        for i in range(30):
            try:
                self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self._socket.connect((self._host, self._port))
                exception = None
                break

            except:
                self._socket = None
                exception = sys.exc_info()

            # Wait half sec and try again
            time.sleep(1)

        if self._socket is None:
            value = ""
            try:
                value = str(exception[1])
            except:
                pass

            raise PeachException("TCP onnection attempt failed: %s" % value)

        self.buff = ""
        self.pos = 0

    def close(self):
        """
        Close connection if open.
        """
        try:
            if self._socket is not None:
                self._socket.close()
        finally:
            self._socket = None

        self.buff = ""
        self.pos = 0

    def send(self, data):
        """
        Send data via sendall.

        @type	data: string
        @param	data: Data to send
        """

        if Peach.Engine.engine.Engine.debug:
            print(">>>>>>>>>>>>>>>>")
            print("tcp.Tcp.send():")
            printHex(data)

        try:
            self._socket.sendall(data)
        except:
            if Peach.Engine.engine.Engine.debug:
                print("Tcp: Sendall failed: " + str(sys.exc_info()[1]))
            raise PublisherSoftException("sendall failed: " + str(sys.exc_info()[1]))

    def _receiveBySize(self, size):
        """
        This is now a buffered receiver.

        @rtype: string
        @return: received data.
        """

        # Do we already a have it?
        if size + self.pos < len(self.buff):
            ret = self.buff[self.pos:self.pos + size]
            self.pos += size
            return ret

        # Only ask for the diff of what we don't already have
        diffSize = (self.pos + size) - len(self.buff)

        try:
            if Peach.Engine.engine.Engine.debug:
                print("Asking for %d, need %d, have %d" % (size, diffSize, len(self.buff) - self.pos))

            self._socket.settimeout(self._timeout)
            ret = self._socket.recv(diffSize)

            if not ret:
                # Socket was closed
                if Peach.Engine.engine.Engine.debug:
                    print("Socket is closed")

                raise PublisherSoftException("Socket is closed")

            if Peach.Engine.engine.Engine.debug:
                print("<<<<<<<<<<<<<<<<<")
                print("tcp.Tcp.receive():")
                printHex(ret)

            self.buff += ret

        except socket.error as e:
            if str(e).find('The socket operation could not complete without blocking') != -1:
                if Peach.Engine.engine.Engine.debug:
                    print("timed out waiting for data")

                raise Timeout(
                    "Timed out waiting for data [%d:%d:%d:%d]" % (len(self.buff), (size + self.pos), size, diffSize))

            elif str(e).find('An existing connection was forcibly') != -1:
                if Peach.Engine.engine.Engine.debug:
                    print("Socket was closed!")

                raise PublisherSoftException("Socket is closed")

            else:
                if Peach.Engine.engine.Engine.debug:
                    print("recv failed: " + str(sys.exc_info()[1]))

                raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1]))

        finally:
            self._socket.settimeout(None)

        ret = self.buff[self.pos:]
        self.pos = len(self.buff)
        return ret

    def _receiveByAvailable(self):
        """
        Receive as much as possible prior to timeout.

        @rtype: string
        @return: received data.
        """
        self._socket.settimeout(self._timeout)

        try:
            ret = self._socket.recv(4096)

            if not ret:
                raise PublisherSoftException("Socket is closed")

            if Peach.Engine.engine.Engine.debug:
                print("<<<<<<<<<<<<<<<<<")
                print("tcp.Tcp.receive():")
                printHex(ret)

            self.buff += ret

        except socket.error as e:
            if str(e).find('The socket operation could not complete without blocking') == -1:
                pass

            else:
                raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1]))

        self._socket.settimeout(None)

        ret = self.buff[self.pos:]
        self.pos = len(self.buff)
        return ret

    def receive(self, size=None):

        if size is None:
            return self._receiveByAvailable()
        else:
            return self._receiveBySize(size)


class TcpListener(Tcp):
    """
    A TCP Listener publisher.  This publisher
    supports the following state actions:

     * start - Start listening
     * stop - Stop listening
     * accept - Accept a client connection
     * close - Close a client connection
    """

    def __init__(self, host, port, timeout=0.25):
        Tcp.__init__(self, host, port, timeout)

        self._listen = None
        self._clientAddr = None

    def start(self):
        self.close()

        if self._listen is None:
            for i in range(3):
                try:
                    self._listen = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    self._listen.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                    self._listen.bind((self._host, self._port))
                    self._listen.listen(1)
                    exception = None
                    break

                except:
                    self._listen = None
                    exception = sys.exc_info()

                time.sleep(0.5)

            if self._listen is None:
                value = ""
                try:
                    value = str(exception[1])
                except:
                    pass

                raise PeachException("TCP bind attempt failed: %s" % value)

        self.buff = ""
        self.pos = 0

    def stop(self):
        try:
            if self._socket is not None:
                self._socket.close()
        except:
            pass

        finally:
            self._socket = None

        # Avoid TIME_WAIT state by not closing our listener

        #try:
        #	if self._listen != None:
        #		self._listen.close()
        #except:
        #	pass
        #
        #finally:
        #	self._listen = None


    def accept(self):
        self.buff = ""
        self.pos = 0

        conn, addr = self._listen.accept()
        self._socket = conn
        self._clientAddr = addr

    def close(self):
        try:
            if self._socket is not None:
                self._socket.close()
        except:
            pass
        finally:
            self._socket = None

    def connect(self):
        raise PeachException("Action 'connect' not supported")


try:
    import win32gui, win32con, win32process
    import sys, time, os, signal

    class TcpListenerLaunchGui(TcpListener):
        """
        Does TcpListener goodness and also can laun a program.  After
        some defined amount of time we will try and close the
        GUI application by sending WM_CLOSE than kill it.
        """

        def __init__(self, host, port, windowname, timeout=0.25):
            TcpListener.__init__(self, host, port, timeout)
            self._windowName = windowname

        def stop(self):
            self.closeApp(self._windowName)
            TcpListener.stop(self)

        def call(self, method, args):
            """
            Launch program to consume file

            @type	method: string
            @param	method: Command to execute
            @type	args: array of objects
            @param	args: Arguments to pass
            """

            realArgs = [method, "/c", method]
            for a in args:
                realArgs.append(a)

            #print "Spawning %s" % method
            os.spawnv(os.P_NOWAIT, os.path.join(os.getenv('SystemRoot'), 'system32', 'cmd.exe'), realArgs)


        @staticmethod
        def enumCallback(hwnd, windowName):
            """
            Will get called by win32gui.EnumWindows, once for each
            top level application window.
            """

            try:

                # Get window title
                title = win32gui.GetWindowText(hwnd)

                # Is this our guy?
                if title.find(windowName) == -1:
                    return

                (threadId, processId) = win32process.GetWindowThreadProcessId(hwnd)

                # Send WM_CLOSE message
                try:
                    win32gui.PostMessage(hwnd, win32con.WM_CLOSE, 0, 0)
                    win32gui.PostQuitMessage(hwnd)
                except:
                    pass

                # Give it upto 5 sec
                for i in range(100):
                    if win32process.GetExitCodeProcess(processId) != win32con.STILL_ACTIVE:
                        # Process exited already
                        return

                    time.sleep(0.25)

                try:
                    # Kill application
                    win32process.TerminateProcess(processId, 0)
                except:
                    pass
            except:
                pass

        def closeApp(self, title):
            """
            Close Application by window title
            """
            #print "CloseApp: %s" % title
            win32gui.EnumWindows(TcpListenerLaunchGui.enumCallback, title)

except:
    pass


class TcpProxyB(TcpListener):
    def __init__(self, host, port):
        TcpListener.__init__(self, host, port)


class TcpProxyA(Tcp):
    def __init__(self, host, port):
        Tcp.__init__(self, host, port)