# The contents of this file are subject to the BitTorrent Open Source License
# Version 1.1 (the License).  You may not copy or use this file, in either
# source code or executable form, except in compliance with the License.  You
# may obtain a copy of the License at http://www.bittorrent.com/license/.
#
# Software distributed under the License is distributed on an AS IS basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
# for the specific language governing rights and limitations under the
# License.

# Written by Bram Cohen

from cStringIO import StringIO
from socket import error as socketerror

protocol_name = 'BitTorrent protocol'

# header, reserved, download id, my id, [length, message]

from twisted.internet.protocol import Protocol, ClientFactory
from twisted.internet import reactor
from twisted.python import log

class NatCheck(object):

    def __init__(self, resultfunc, downloadid, peerid, ip, port):
        self.resultfunc = resultfunc
        self.downloadid = downloadid
        self.peerid = peerid
        self.ip = ip
        self.port = port
        self.answered = False

        factory = NatCheckProtocolFactory(self, downloadid, peerid)

        reactor.connectTCP(ip, port, factory)

    def answer(self, result):
        if not self.answered:
            self.answered = True
            log.msg('NAT check for %s:%i is %s' % (self.ip, self.port, result))
            self.resultfunc(result, self.downloadid, self.peerid, self.ip, self.port)

class NatCheckProtocolFactory(ClientFactory):
    def __init__(self, natcheck, downloadid, peerid):
        self.natcheck = natcheck
        self.downloadid = downloadid
        self.peerid = peerid

    def startedConnecting(self, connector):
        log.msg('Started to connect.')
        
    def buildProtocol(self, addr):
        return NatCheckProtocol(self, self.downloadid, self.peerid)
    
    def clientConnectionLost(self, connector, reason):
        self.natcheck.answer(False)
        log.msg('Lost connection.  Reason: %s' % reason)
    
    def clientConnectionFailed(self, connector, reason):
        self.natcheck.answer(False)
        log.msg('Connection failed. Reason: %s' % reason)

class NatCheckProtocol(Protocol):
    def __init__(self, factory, downloadid, peerid):
        self.factory = factory
        self.downloadid = downloadid
        self.peerid = peerid
        self.data = ''
        self.received_protocol_name_len = None
        self.received_protocol_name = None
        self.received_reserved = None
        self.received_downloadid = None
        self.received_peerid = None

    def connectionMade(self):
        self.transport.write(chr(len(protocol_name)))
        self.transport.write(protocol_name)
        self.transport.write(chr(0) * 8)
        self.transport.write(self.downloadid)

    def dataReceived(self, data):
        self.data += data

        if self.received_protocol_name_len is None:
            if len(self.data) >= 1:
                self.received_protocol_name_len = ord(self.data[0])
                self.data = self.data[1:]
                if self.received_protocol_name_len != len(protocol_name):
                    self.factory.natcheck.answer(False)
                    self.transport.loseConnection()
                    return
            else:
                return

        if self.received_protocol_name is None:
            if len(self.data) >= self.received_protocol_name_len:
                self.received_protocol_name = self.data[:self.received_protocol_name_len]
                self.data = self.data[self.received_protocol_name_len:]
                if self.received_protocol_name != protocol_name:
                    log.err('Received protocol name did not match!')
                    self.factory.natcheck.answer(False)
                    self.transport.loseConnection()
                    return
            else:
                return

        if self.received_reserved is None:
            if len(self.data) >= 8:
                self.received_reserved = self.data[:8]
                self.data = self.data[8:]
            else:
                return

        if self.received_downloadid is None:
            if len(self.data) >= 20:
                self.received_downloadid = self.data[:20]
                self.data = self.data[20:]
                if self.received_downloadid != self.downloadid:
                    log.err('Received download id did not match!')
                    self.factory.natcheck.answer(False)
                    self.transport.loseConnection()
                    return
            else:
                return

        if self.received_peerid is None:
            if len(self.data) >= 20:
                log.msg('Peerid length: %i' % len(self.peerid))
                self.received_peerid = self.data[:20]
                self.data = self.data[20:]
                log.msg('Received: %s' % self.received_peerid.encode('hex'))
                log.msg('Received: %s' % self.received_peerid.encode('quoted-printable'))
                log.msg('Expected: %s' % self.peerid.encode('hex'))
                log.msg('Expected: %s' % self.peerid.encode('quoted-printable'))
                if self.received_peerid != self.peerid:
                    log.err('Received peer id did not match!')
                    self.factory.natcheck.answer(False)
                    self.transport.loseConnection()
                    return
            else:
                return

        if self.received_protocol_name == protocol_name and self.received_downloadid == self.downloadid and self.received_peerid == self.peerid:
            self.factory.natcheck.answer(True)
            self.transport.loseConnection()