from twisted.internet.interfaces import ITransport
from twisted.internet.base      import BaseConnector
from twisted.internet           import reactor, tcp
from twisted.internet           import protocol, defer
from twisted.python             import failure
import struct, re, socket, sys

from sockserror import *

from debug import LogEvent, INFO, WARN, ERROR

# Used to distinguish IP address from domain name
# TODO: should this be optimized somehow?
#
_ip_regex = re.compile ("\d\d?\d?\.\d\d?\d?\.\d\d?\d?\.\d\d?\d?")


class ClientProtocol (protocol.Protocol):
    """ This protocol that talks to SOCKS5 server from client side.
    """
    __implements__ = ITransport,
    disconnecting = 0

    def __init__(self, sockshost, socksport, host, port, factory, otherProtocol,
        method="CONNECT", login=None, password=None):
        """ Initializes SOCKS session
        
        @type sockshost: string
        @param sockshost: Domain name or ip address of intermediate SOCKS 
server.

        @type socksport: int
        @param socksport: Port number of intermediate server.

        @type host: string
        @param host: Domain name or ip address where should connect or bind.

        @type port: int
        @param port: Port number where to connect or bind.

        @type otherProtocol: object
        @param otherProtocol: Initialised protocol instance, which will receive
            all I/O and events after SOCKS connected.

        @type login: string
        @param login: Sets user name if SOCKS server requires us to
            authenticate.

        @type password: string
        @param password: Sets user password if SOCKS server requires us
            to authenticate.

        @type method: string
        @param method: What to do: may be \"CONNECT\" only. Other
            methods are currently unsupported.
        """
        # login and password are limited to 256 chars
        #
        if login is not None and len (login) > 255:
            raise LoginTooLongError()

        if password is not None and len (password) > 255:
            raise PasswordTooLongError()

        # save information
        #
        self.method         = method
        self.host           = host
        self.port           = port
        self.login          = login
        self.password       = password
        self.state          = "mustNotReceiveData"
        self.otherProtocol  = otherProtocol
        self.factory        = factory
	    
    def connectionMade(self):
        # prepare connection string with available authentication methods
        #
	
        LogEvent(INFO, msg="SOCKS5.connectionMade")
        methods = "\x00"
        if not self.login is None: methods += "\x02"

        connstring = struct.pack ("!BB", 5, len (methods))

        self.transport.write (connstring + methods)
        self.state = "gotHelloReply"

    def dataReceived (self, data):
        LogEvent(INFO, msg="SOCKS state=" + self.state)
        method = getattr(self, 'socks_%s' % (self.state), 
            self.socks_thisMustNeverHappen)
        method (data)

    def socks_thisMustNeverHappen (self, data):
        self.transport.loseConnection()
        raise UnhandledStateError ("This SOCKS5 self.state (%s) "\
            "must never happen %s" % (self.state, self))

    def socks_mustNotReceiveData (self, data):
        """ This error might occur when server tells something into connection
        right after connection is established. Server in this case is
        certainly not SOCKS.
        """
        self.transport.loseConnection()
        self.factory.clientConnectionFailed (self, failure.Failure (
            UnexpectedDataError ("Server must not send data before client %s" % 
self)))

    def socks_gotHelloReply (self, data):
        """ Receive server greeting and send authentication or ask to
        execute requested method right now.
        """
        if data == "\x05\xFF":
            # No acceptable methods. We MUST close
            #
	    LogEvent(WARN, msg="No acceptable methods, closing connection")
            self.transport.loseConnection()
            return

        elif data == "\x05\x00":
            # Anonymous access allowed - let's issue connect
            #
            self.sendCurrentMethod()

        elif data == "\x05\x02":
            # Authentication required
            #
            self.sendAuth()

        else:
            self.transport.loseConnection()
            self.factory.clientConnectionFailed (self, failure.Failure (
                UnhandledData ("Server returned unknown reply in gotHelloReply")))

        # From now on SOCKS server considered alive - we've got reply
        #
        self.factory.status = "connected"

    def socks_gotAuthReply (self, data):
        """ Called when client received server authentication reply,
            we or close connection or issue "CONNECT" command
        """
        if data == "\x05\x00":
            self.sendCurrentMethod()

    def sendAuth (self):
        """ Prepare login/password pair and send it to the server
        """
        command = "\x05%s%s%s%s" % (chr (len (self.login)), self.login,
            chr (len (self.password)), self.password)
        self.transport.write (command)

        self.state = "gotAuthReply"

    def sendCurrentMethod (self):
        method = getattr(self, 'socks_method_%s' % (self.method), 
            self.socks_method_UNKNOWNMETHOD)
        method()

    def socks_method_UNKNOWNMETHOD (self):
        self.transport.loseConnection()
        self.factory.clientConnectionFailed (self, failure.Failure (
            UnknownMethod ("Method %s is unknown %s" % (self.method, self))))

    def socks_method_CONNECT (self):
        # Check if we have ip address or domain name
        #
	LogEvent(INFO, msg="socks_method_CONNECT host = " + self.host)

 	# The FaceTime SOCKS5 proxy treats IP addr the same way as hostname
       # if _ip_regex.match (self.host):
       #     # we have dotted quad IP address
       #     addressType = 1
       #     address = socket.inet_aton (self.host)
       # else:
       #     # we have host name
       #     address = self.host
       #     addressType = 3

	address = self.host
	addressType = 3
	addressLen = len(address)

        #Protocol version=5, Command=1 (CONNECT), Reserved=0
        #command = struct.pack ("!BBBB", 5, 1, 0, addressType)

	command = struct.pack ("!BBBBB", 5, 1, 0, addressType,addressLen)
        portstr = struct.pack ("!H", self.port)

        self.transport.write (command + address + portstr)
        self.state = "gotConnectReply"

    def socks_gotConnectReply (self, data):
        """ Called after server accepts or rejects CONNECT method.
        """
        if data[:2] == "\x05\x00":
            # No need to analyze other fields of reply, we are done
            #
            self.state = "done"
            self.factory.status = "established"

            self.otherProtocol.transport = self
            self.otherProtocol.connectionMade()
            return

        errcode = ord (data[1])

        if errcode < len (SOCKS_errors):
	    LogEvent(ERROR, msg='SOCKS error: %s' % SOCKS_errors[errcode])
            self.transport.loseConnection()
            self.factory.clientConnectionFailed (self, failure.Failure (
                ConnectError ("%s %s" % (SOCKS_errors[errcode], self))))
        else:
            self.transport.loseConnection()
            self.factory.clientConnectionFailed (self, failure.Failure (
                ConnectError ("Unknown SOCKS error after CONNECT request issued %s" % (self))))

    def socks_done (self, data):
        """ Proxy received data to other protocol.
        """
        self.otherProtocol.dataReceived (data)
    #
    # Transport relaying
    #
    def write(self, data):
        self.transport.write(data)

    def writeSequence(self, data):
        self.transport.writeSequence(data)

    def loseConnection(self):
        self.disconnecting = 1
        self.transport.loseConnection()

    def getPeer(self):
        return self.transport.getPeer()

    def getHost(self):
        return self.transport.getHost()
    
    def registerProducer(self, producer, streaming):
        self.transport.registerProducer(producer, streaming)

    def unregisterProducer(self):
        self.transport.unregisterProducer()

    def stopConsuming(self):
        self.transport.stopConsuming()
	
    def setTcpNoDelay(self, enabled):
	pass # as possible implement it? SOCKS 5 protocol have no support for one

class ClientConnector (tcp.Connector):
    """Object used to connect to some host using intermediate server
    supporting SOCKS5 protocol.

    This IConnector manages one connection.
    """
    def __init__(self, sockshost, socksport, host, port, otherFactory,
        reactor=None, method="CONNECT", login=None, password=None,
        timeout=None, readableID=None):
        """ Creates IConnector to connect through SOCKS

        @type sockshost: string
        @param sockshost: SOCKS5 compliant server address.

        @type socksport: int
        @param socksport: Port to use when connecting to SOCKS.

        @type timeout: float
        @param timeout: Time to wait until client connects, then fail.

        @type readableID: string
        @param readableID: Some human readable ID for this connection.

        See ClientProtocol constructor for details on other params.
        """
        factory = ClientFactory (method=method, sockshost=sockshost,
            socksport=socksport, host=host, port=port, login=login,
            password=password, otherFactory=otherFactory, timeout=timeout,
            readableID=readableID)

        tcp.Connector.__init__ (self, host=sockshost, port=socksport,
            factory=factory, timeout=timeout, bindAddress=None,
            reactor=reactor)

class ClientFactory (protocol.ClientFactory):
    def __init__(self, sockshost, socksport, host, port, otherFactory,
        method="CONNECT", login=None, password=None, timeout=60,
        readableID=None):
        """ Factory creates SOCKS5 client protocol to connect through it.
        See ClientProtocol constructor for details on params.
        
        @type globalTimeout: int
        @param globalTimeout: Seconds before connection is completely and
            unconditionally closed as is.

        @type readableID: string
        @param readableID: Some human readable ID for this connection.
        """
        self.sockshost      = sockshost
        self.socksport      = socksport
        self.host           = host
        self.port           = port
        self.method         = method
        self.login          = login
        self.password       = password
        self.otherFactory   = otherFactory
        self.timeout        = timeout
        self.readableID     = readableID

        # This variable contains current status of SOCKS connection,
        # useful for diagnosting connection, without knowing SOCKS
        # internal states. One of: "unconnected", "connected" (SOCKS
        # server replied and is alive), "established"
        #
        self.status = "unconnected"

    def startedConnecting (self, connector):
        # Set global timeout
        #
	if self.timeout is not None:
           LogEvent(INFO, msg="Set timeout %d sec" % self.timeout)
           delayedcall = reactor.callLater (self.timeout, self.onTimeout,
           				    connector)
           setattr (self, "delayed_timeout_call", delayedcall)

        # inherited
        #
        protocol.ClientFactory.startedConnecting (self, connector)

    def onTimeout (self, connector):
        """ Timeout occured, can't continue and should stop immediately
        and unconditionally in the whatever state I am.
        """
        connector.disconnect()
        LogEvent(INFO, msg="%s timeout %d sec" % (self, self.timeout))
        self.clientConnectionFailed (self, failure.Failure (
            GlobalTimeoutError ("Timeout %s" % self)))

    def stopFactory(self):
        """ Do cleanups such as cancelling timeout
        """
        try:
	    if self.timeout is not None:
               self.delayed_timeout_call.cancel()
        except:
            pass

        protocol.ClientFactory.stopFactory (self)

    def buildProtocol (self, a):
        """ Connection is successful, create protocol and let it talk to peer.
        """
        proto = ClientProtocol (sockshost=self.sockshost,
            socksport=self.socksport, host=self.host, port=self.port,
            method=self.method, login=self.login, password=self.password,
            otherProtocol=self.otherFactory.buildProtocol (self.sockshost),
            factory=self)

        proto.factory = self
        return proto

    def __repr__ (self):
        return "<SOCKS %s>" % self.readableID

    def clientConnectionLost(self, connector, reason):
        # If flag indicates that connection may not be lost
        #
        rmap = {"reason": reason, "socks": self.status}

        try:
            if self.status != "established":
                # Tell about error
                #
                LogEvent(ERROR, msg="Connection LOST before SOCKS established %s" % self)
                self.otherFactory.clientConnectionLost (connector, rmap)
            else:
                self.otherFactory.clientConnectionLost (connector, rmap)
        except:
            ei = sys.exc_info()
            if not str (ei[0]).count ("AlreadyCalled"):
                raise
    
    def clientConnectionFailed(self, connector, reason):
        # I can't know where to get deferred, let factory do this itself
        #
        rmap = {"reason": reason, "socks": self.status}

        try:
            if self.status != "established":
                LogEvent(ERROR, msg="Connection FAILED before SOCKS established %s" % self)
		# prevent Traceback in Twisted
                #self.otherFactory.clientConnectionFailed (connector, rmap)
            else:
                self.otherFactory.clientConnectionFailed (connector, rmap)
        except:
            ei = sys.exc_info()
            if not str (ei[0]).count ("AlreadyCalled"):
                raise


class ProxyClientCreator(protocol.ClientCreator):

	def connectSocks5Proxy(self,remotehost,remoteport,proxy,proxyport,id):
	   d = defer.Deferred()
	   f = protocol._InstanceFactory(self.reactor,
	                                self.protocolClass(*self.args, **self.kwargs),d)
	   self.reactor.connectWith(ClientConnector,
			            host=remotehost,port=remoteport,
				    sockshost=proxy,socksport=proxyport,
				    otherFactory=f,
				    readableID=id)
	   return d




