nattraversal.py

bittorrent source by python. please enjoy

# someday: http://files.dns-sd.org/draft-nat-port-mapping.txt
# today: http://www.upnp.org/

import os
import sys
import Queue
import socket
import random
import logging
import urlparse
if os.name == 'nt':
    import pywintypes
    import win32com.client

from BitTorrent import app_name, defer
from BitTorrent.platform import os_version
from BitTorrent.sparse_set import SparseSet
from BitTorrent.RawServer_twisted import RawServer, Handler
from BitTorrent.BeautifulSupe import BeautifulSupe, Tag
from BitTorrent.yielddefer import launch_coroutine, _wrap_task
from BitTorrent.HostIP import get_host_ip, get_deferred_host_ip
from BitTorrent.obsoletepythonsupport import has_set, set
from twisted import internet
import twisted.copyright

from urllib2 import URLError, HTTPError, Request
from httplib import BadStatusLine

#bleh
from urllib import urlopen, FancyURLopener, addinfourl
from httplib import HTTPResponse

import BitTorrent.stackthreading as threading

nat_logger = logging.getLogger('NatTraversal')
nat_logger.setLevel(logging.WARNING)

def UnsupportedWarning(s):
    nat_logger.warning("NAT Traversal warning " + ("(%s: %s)."  % (os_version, s)))

def UPNPError(s):
    nat_logger.error("UPnP ERROR: " + ("(%s: %s)."  % (os_version, s)))

class UPnPException(Exception):
    pass

class NATEventLoop(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.queue = Queue.Queue()
        self.killswitch = defer.DeferredEvent()
        def ignore(*a, **kw):
            pass
        self.killswitch.addCallback(ignore)            
        self.setDaemon(True)

    def run(self):

        while not self.killswitch.isSet():

            (f, a, kw) = self.queue.get()
            
            try:
                nat_logger.debug("NATEventLoop Event: %s" % f.__name__)
                f(*a, **kw)
                nat_logger.debug("NATEventLoop Event: %s finished." % f.__name__)
            except:
                # sys can be none during interpritter shutdown
                if sys is None:
                    break
                nat_logger.exception("Error in NATEventLoop for %s" % str(f.__name__))


class NatTraverser(object):
    def __init__(self, rawserver):
        self.rawserver = rawserver
        
        self.register_requests = []
        self.unregister_requests = []
        self.list_requests = []

        self.service = None
        self.services = []
        self.current_service = 0

        if self.rawserver.config['upnp']:
            if os.name == 'nt':
                self.services.append(WindowsUPnP)
            self.services.append(ManualUPnP)

        self.event_loop = NATEventLoop()
        self.event_loop.start()
        
        self.resume_init_services()

    def add_task(self, f, *a, **kw):
        self.event_loop.queue.put((f, a, kw))        

    def init_services(self):
        # this loop is a little funny so a service can resume the init if it fails later
        if not self.rawserver.config['upnp']:
            return
        while self.current_service < len(self.services):
            service = self.services[self.current_service]
            self.current_service += 1
            try:
                nat_logger.info("Trying: %s" % service.__name__)
                service(self)
                break
            except Exception, e:
                nat_logger.warning(unicode(e.args[0]))
        else:
            e = "Unable to detect any UPnP services"
            UnsupportedWarning(e)
            self._cancel_queue(e)

    def resume_init_services(self):
        self.add_task(self.init_services)

    def attach_service(self, service):
        nat_logger.info("Using: %s" % type(service).__name__)
        self.service = service
        self.add_task(self._flush_queue)

    def detach_service(self, service):
        if service != self.service:
            nat_logger.error("Service: %s is not in use!" % type(service).__name__)
            return
        nat_logger.info("Detached: %s" % type(service).__name__)
        self.service = None
        
    def _flush_queue(self):
        if self.service:
            for mapping in self.register_requests:
                self.add_task(self.service.safe_register_port, mapping)
            self.register_requests = []
        
            for request in self.unregister_requests:
                # unregisters can block, because they occur at shutdown
                self.service.unregister_port(*request)
            self.unregister_requests = []

            for request in self.list_requests:
                self.add_task(self._list_ports, request)
            self.list_requests = []

    def _cancel_queue(self, e):
        for mapping in self.register_requests:
            mapping.d.errback(e)
        self.register_requests = []

        # can't run or cancel blocking removes    
        self.unregister_requests = []

        for request in self.list_requests:
            request.errback(e)
        self.list_requests = []

    def _gen_deferred(self):
        return defer.ThreadableDeferred(_wrap_task(self.rawserver.external_add_task))

    def register_port(self, external_port, internal_port, protocol,
                      host = None, service_name = None, remote_host=''):
        mapping = UPnPPortMapping(external_port, internal_port, protocol,
                                  host, service_name, remote_host)
        mapping.d = self._gen_deferred()
        self.register_requests.append(mapping)

        self.add_task(self._flush_queue)

        return mapping.d
    
    def unregister_port(self, external_port, protocol):
        self.unregister_requests.append((external_port, protocol))

        # unregisters can block, because they occur at shutdown
        self._flush_queue()

    def _list_ports(self, d):
        matches = self.service._list_ports()
        d.callback(matches)

    def list_ports(self):
        d = self._gen_deferred()
        self.list_requests.append(d)
        self.add_task(self._flush_queue)
        return d
        

        

class NATBase(object):

    def safe_register_port(self, new_mapping):

        # check for the host now, while we're in the thread and before
        # we need to read it.
        new_mapping.populate_host()
        
        nat_logger.info("You asked for: " + str(new_mapping))
        new_mapping.original_external_port = new_mapping.external_port
        mappings = self._list_ports()

        used_ports = []
        for mapping in mappings:
            # only consider ports which match the same protocol
            if mapping.protocol == new_mapping.protocol:
                # look for exact matches
                if (mapping.host == new_mapping.host and
                    mapping.internal_port == new_mapping.internal_port):
                    # the service name could not match, that's ok.
                    new_mapping.d.callback(mapping.external_port)
                    nat_logger.info("Already effectively mapped: " + str(new_mapping))
                    return 
                # otherwise, add it to the list of used external ports
                used_ports.append(mapping.external_port)

        used_ports.sort()
        used_ports = SparseSet(used_ports)

        all_ports = SparseSet()
        all_ports.add(1024, 65535)
        free_ports = all_ports - used_ports
        new_mapping.external_port = random.choice(free_ports)

        nat_logger.info("I'll give you: " + str(new_mapping))
        self.register_port(new_mapping)
        
    def register_port(self, port):
        pass
    def unregister_port(self, external_port, protocol):
        pass
    def _list_ports(self):
        pass

class UPnPPortMapping(object):
    def __init__(self, external_port, internal_port, protocol,
                 host = None, service_name = None, remote_host=''):
        self.external_port = int(external_port)
        self.internal_port = int(internal_port)
        self.protocol = protocol

        self.host = host
        self.remote_host = ''

        if service_name:
            self.service_name = service_name
        else:
            self.service_name = app_name

        self.d = defer.Deferred()

    def populate_host(self):
        # throw out '' or None or ints, also look for semi-valid IPs
        if not isinstance(self.host, str) or self.host.count('.') < 3:
            self.host = get_host_ip()
        
    def __str__(self):
        if not self.remote_host:
            remote = 'external'
        else:
            remote = self.remote_host
        return "%s %s %s:%d %s:%d" % (self.service_name, self.protocol,
                                      self.remote_host,
                                      self.external_port,
                                      self.host, self.internal_port)

def VerifySOAPResponse(request, response):
    if response.code != 200:
        raise HTTPError(request.get_full_url(),
                        response.code, str(response.msg) + " (unexpected SOAP response code)",
                        response.info(), response)

    data = response.read()
    bs = BeautifulSupe(data)
    # On Matt's Linksys WRT54G rev 4 v.1.0 I saw u: instead of m:
    # and ignoring that caused the router to crash
    soap_response = bs.scour("m:", "Response")
    if not soap_response:
        raise HTTPError(request.get_full_url(),
                        response.code, str(response.msg) +
                        " (incorrect SOAP response method)",
                        response.info(), response)
    return soap_response[0]
    
def SOAPResponseToDict(soap_response):
    result = {}
    for tag in soap_response.child_elements():
        value = None
        if tag.contents:
            value = str(tag.contents[0])
        result[tag.name] = value
    return result

def SOAPErrorToString(response):
    if not isinstance(response, Exception):
        data = response.read()
        bs = BeautifulSupe(data)
        error = bs.first('errorDescription')
        if error:
            return str(error.contents[0])
    return str(response)

_urlopener = None
def urlopen_custom(req, rawserver):
    global _urlopener

    if not _urlopener:
        opener = FancyURLopener()        
        _urlopener = opener
        #remove User-Agent
        del _urlopener.addheaders[:]

    if not isinstance(req, str):
        #for header in r.headers:
        #    _urlopener.addheaders.append((header, r.headers[header]))
        #return _urlopener.open(r.get_full_url(), r.data)
        
        # All this has to be done manually, since httplib and urllib 1 and 2
        # add headers to the request that some routers do not accept.
        # A minimal, functional request includes the headers:
        # Content-Length
        # Soapaction
        # I have found the following to be specifically disallowed:
        # User-agent
        # Connection
        # Accept-encoding

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        (scheme, netloc, path, params, query, fragment) = urlparse.urlparse(req.get_full_url())

        if not scheme.startswith("http"):
            raise ValueError("UPnP URL scheme is not http: " + req.get_full_url())

        if len(path) == 0:
            path = '/'

        if netloc.count(":") > 0:
            host, port = netloc.split(':', 1)
            try:
                port = int(port)
            except:
                raise ValueError("UPnP URL port is not int: " + req.get_full_url())
        else:
            host = netloc
            port = 80
        
        header_str = ''
        data = ''
        method = ''
        header_str = " " + path + " HTTP/1.0\r\n"
        if req.has_data():
            method = 'POST'
            header_str = method + header_str
            header_str += "Content-Length: " + str(len(req.data)) + "\r\n"
            data = req.data + "\r\n"
        else:
            method = 'GET'
            header_str = method + header_str
            
        header_str += "Host: " + host + ":" + str(port) + "\r\n"
        
        for header in req.headers:
            header_str += header + ": " + str(req.headers[header]) + "\r\n"

        header_str += "\r\n"
        data = header_str + data

        try:
            rawserver.add_pending_connection(host)
            s.connect((host, port))
        finally:
            rawserver.remove_pending_connection(host)
            
        s.send(data)
        r = HTTPResponse(s, method=method)
        r.begin()

        r.recv = r.read
        fp = socket._fileobject(r)

        resp = addinfourl(fp, r.msg, req.get_full_url())
        resp.code = r.status
        resp.msg = r.reason
                   
        return resp

    return _urlopener.open(req)


class ManualUPnP(NATBase, Handler):

    upnp_addr = ('239.255.255.250', 1900)

    search_string = ('M-SEARCH * HTTP/1.1\r\n' +
                     'Host:239.255.255.250:1900\r\n' +