# Written by Bram Cohen, Jie Yang
# see LICENSE.txt for license information
import sys
from cStringIO import StringIO
from binascii import b2a_hex
from socket import error
from urllib import quote
from struct import unpack
from sha import sha
from time import time,ctime
from __init__ import *
from BitTornado.BT1.Encrypter import protocol_name,option_pattern
try:
True
except:
True = 1
False = 0
DEBUG = False
MAX_INCOMPLETE = 8
def toint(s):
return long(b2a_hex(s), 16)
def make_readable(s):
if not s:
return ''
if quote(s).find('%') >= 0:
return b2a_hex(s).upper()
return '"'+s+'"'
class IncompleteCounter:
def __init__(self):
self.c = 0
def increment(self):
self.c += 1
def decrement(self):
self.c -= 1
def toomany(self):
return self.c >= MAX_INCOMPLETE
incompletecounter = IncompleteCounter()
# header, reserved, download id, my id, [length, message]
class Connection: # OverlaySocket, a better name for it?
def __init__(self, Encoder, singsock, id, ext_handshake = False, dns = None):
self.Encoder = Encoder
self.singsock = singsock # SocketHandler.SingleSocket
self.connecter_conn = None # OverlayConnecter.Connection
self.connecter = Encoder.connecter
self.id = id
self.dns = dns
self.readable_id = make_readable(id)
self.locally_initiated = (id != None)
self.complete = False
self.keepalive = lambda: None
self.closed = False
self.buffer = StringIO()
self.overlay_version = CurrentVersion
if self.locally_initiated:
incompletecounter.increment()
if self.locally_initiated or ext_handshake:
#if DEBUG:
# print >> sys.stderr,"olencoder: writing header"
self.singsock.write(chr(len(protocol_name)) + protocol_name +
option_pattern + self.Encoder.download_id)
if ext_handshake:
self.Encoder.connecter.external_connection_made += 1
self.singsock.write(self.Encoder.my_id)
self.next_len, self.next_func = 20, self.read_peer_id
else:
self.next_len, self.next_func = 1, self.read_header_len
#TODO: auto close management
#self.Encoder.raw_server.add_task(self._auto_close, 15)
def get_ip(self, real=False):
return self.singsock.get_ip(real)
def get_port(self, real=False):
return self.singsock.get_port(real)
def get_dns(self):
if not self.dns:
self.dns = (self.get_ip(True), self.get_port(True))
return self.dns
def get_myip(self, real=False):
return self.singsock.get_myip(real)
def get_myport(self, real=False):
return self.singsock.get_myport(real)
def get_id(self):
return self.id
def get_readable_id(self):
return self.readable_id
def is_locally_initiated(self):
return self.locally_initiated
def is_flushed(self):
return self.singsock.is_flushed()
def set_options(self, s):
pass # used by future
def read_header_len(self, s):
if ord(s) != len(protocol_name):
return None
return len(protocol_name), self.read_header
def read_header(self, s):
if s != protocol_name:
return None
return 8, self.read_reserved
def read_reserved(self, s):
if DEBUG:
print >> sys.stderr,"olencoder: Reserved bits:", `s`
self.set_options(s)
return 20, self.read_download_id
def read_download_id(self, s):
if s != self.Encoder.download_id:
return None
if not self.locally_initiated:
self.Encoder.connecter.external_connection_made += 1
self.singsock.write(chr(len(protocol_name)) + protocol_name +
option_pattern + self.Encoder.download_id + self.Encoder.my_id)
return 20, self.read_peer_id
def read_peer_id(self, s):
if not self.id:
self.id = s
self.readable_id = make_readable(s)
else:
if s != self.id:
return None
self.complete = self.Encoder.got_id(self)
if not self.complete:
return None
if not self.version_supported(s[16:18], s[18:20]): # versioning protocol
if DEBUG:
print >> sys.stderr,"olencoder: We don't support peer's version of the protocol"
return None
if self.locally_initiated:
self.singsock.write(self.Encoder.my_id)
incompletecounter.decrement()
c = self.Encoder.connecter.connection_made(self, self.dns)
self.keepalive = c.send_keepalive
self.connecter_conn = c
def notify(connection=c):
self.Encoder.overlayswarm.connectionMade(connection)
self.Encoder.raw_server.add_task(notify, 0)
return 4, self.read_len
def version_supported(self, low_ver_str, cur_ver_str):
""" overlay swarm versioning solution: use last 4 bytes in PeerID """
low_ver = unpack('<H', low_ver_str)[0]
cur_ver = unpack('<H', cur_ver_str)[0]
if cur_ver != CurrentVersion:
if low_ver > CurrentVersion: # the other's version is too high
return False
if cur_ver < LowestVersion: # the other's version is too low
return False
if cur_ver < CurrentVersion and \
cur_ver not in SupportedVersions: # the other's version is not supported
return False
if cur_ver < CurrentVersion: # set low version as our version
self.overlay_version = cur_ver
return True
def read_len(self, s):
l = toint(s)
if l > self.Encoder.max_len:
return None
return l, self.read_message
def read_message(self, s):
if s != '':
self.connecter.got_message(self, s)
return 4, self.read_len
def read_dead(self, s):
return None
def _auto_close(self):
if not self.complete:
self.close()
def close(self):
if not self.closed:
self.singsock.close()
self.sever()
def sever(self):
if self.closed:
return
self.closed = True
if self.Encoder.connections.has_key(self.singsock):
del self.Encoder.connections[self.singsock]
if self.complete and self.connecter_conn:
self.connecter.connection_lost(self)
elif self.locally_initiated:
incompletecounter.decrement()
if self.connecter_conn:
self.Encoder.overlayswarm.connectionLost(self.connecter_conn)
def send_message_raw(self, message):
if not self.closed:
self.singsock.write(message)
def data_came_in(self, connection, s):
self.Encoder.measurefunc(len(s))
while 1:
if self.closed:
return
i = self.next_len - self.buffer.tell()
if i > len(s):
self.buffer.write(s)
return
self.buffer.write(s[:i])
s = s[i:]
m = self.buffer.getvalue()
self.buffer.reset()
self.buffer.truncate()
try:
x = self.next_func(m)
except:
self.next_len, self.next_func = 1, self.read_dead
raise
if x is None:
self.close()
return
self.next_len, self.next_func = x
def connection_lost(self, singsock):
if DEBUG:
print >> sys.stderr,"olencoder: connection_lost", self.dns, ctime(time())
if self.Encoder.connections.has_key(singsock):
self.sever()
def connection_flushed(self, singsock):
if self.complete:
self.connecter.connection_flushed(self)
class OverlayEncoder:
def __init__(self, overlayswarm, connecter, raw_server, my_id, max_len,
schedulefunc, keepalive_delay, download_id,
measurefunc, config):
self.overlayswarm = overlayswarm
self.connecter = connecter
self.raw_server = raw_server
self.my_id = my_id
self.max_len = max_len
self.schedulefunc = schedulefunc
self.keepalive_delay = keepalive_delay
self.download_id = download_id
self.measurefunc = measurefunc
self.config = config
self.connections = {}
self.banned = {}
self.to_connect = []
self.paused = False
if self.config['max_connections'] == 0:
self.max_connections = 2 ** 30
else:
self.max_connections = self.config['max_connections']
# overlay swarm doesn't send keeplive if not necessary
#schedulefunc(self.send_keepalives, keepalive_delay)
def send_keepalives(self):
if self.paused:
return
for c in self.connections.values():
c.keepalive()
def keep_alive(self):
self.send_keepalives()
self.schedulefunc(self.keep_alive, self.keepalive_delay)
def start_connections(self, list):
if not self.to_connect:
self.raw_server.add_task(self._start_connection_from_queue)
self.to_connect = list
def _start_connection_from_queue(self):
if self.connecter.external_connection_made:
max_initiate = self.config['max_initiate']
else:
max_initiate = int(self.config['max_initiate']*1.5)
cons = len(self.connections)
if cons >= self.max_connections or cons >= max_initiate:
delay = 60
elif self.paused or incompletecounter.toomany():
delay = 1
else:
delay = 0
dns, id = self.to_connect.pop(0)
self.start_connection(dns, id)
if self.to_connect:
self.raw_server.add_task(self._start_connection_from_queue, delay)
def start_connection(self, dns, id):
if (self.paused
or len(self.connections) >= self.max_connections
or id == self.my_id
or self.banned.has_key(dns[0])):
return True
## Arno: there is something wrong here, SecureOverlay thinks there is
## no connection, but this encoder still thinks there is sometimes.
## FIXME
#for v in self.connections.values(): # avoid duplicated connectiion from a single ip
# if v is None:
# continue
# if id and v.id == id:
# return True
# ip = v.get_ip(True)
# if ip != 'unknown' and ip == dns[0]: # forbid to setup multiple connection on overlay swarm
# if DEBUG:
# print >> sys.stderr,"olencoder: Using existing connection to",dns
# return True
try:
if DEBUG:
print >> sys.stderr,"olencoder: Setting up new connection to", dns, ctime(time())
singsock = self.raw_server.start_connection(dns)
con = Connection(self, singsock, id, dns = dns)
self.connections[singsock] = con
singsock.set_handler(con)
except socketerror:
return False
return True
def _start_connection(self, dns, id):
def foo(self=self, dns=dns, id=id):
self.start_connection(dns, id)
self.schedulefunc(foo, 0)
def got_id(self, connection):
if connection.id == self.my_id:
self.connecter.external_connection_made -= 1
return False
ip = connection.get_ip(True)
if self.config['security'] and self.banned.has_key(ip):
return False
for v in self.connections.values():
if connection is not v:
if connection.id == v.id:
return False
if self.config['security'] and ip != 'unknown' and ip == v.get_ip(True):
v.close()
return True
def external_connection_made(self, connection):
""" Remotely initiated connection """
if self.paused or len(self.connections) >= self.max_connections:
connection.close()
return False
con = Connection(self, connection, None)
self.connections[connection] = con
connection.set_handler(con)
return True
def externally_handshaked_connection_made(self, connection, options, msg_remainder):
if self.paused or len(self.connections) >= self.max_connections:
connection.close()
return False
con = Connection(self, connection, None, True)
self.connections[connection] = con
connection.set_handler(con)
if msg_remainder:
con.data_came_in(con, msg_remainder)
return True
def close_all(self):
for c in self.connections.values():
c.close()
self.connections = {}
def ban(self, ip):
self.banned[ip] = 1
def pause(self, flag):
self.paused = flag
|