instrumentation/next-share/BaseLib/Test/test_na_extend_hs.py

   1 # Written by Arno Bakker
   2 # see LICENSE.txt for license information
   3 #
   4 # TODO: Let Tribler initiate a BT connection to us. We then pretend to be old client
   5 # and then he should iniate an OL connection to us.
   6 #
   7
   8 import unittest
   9 import sys
  10 import time
  11 import socket
  12 from traceback import print_exc
  13 from types import DictType,IntType
  14
  15 from BaseLib.Test.test_extend_hs import TestExtendHandshake
  16 from btconn import BTConnection
  17 from BaseLib.Core.BitTornado.bencode import bencode,bdecode
  18 from BaseLib.Core.BitTornado.BT1.MessageID import *
  19 from BaseLib.Core.BitTornado.BT1.track import compact_ip
  20 from BaseLib.Core.Utilities.Crypto import sha
  21
  22 DEBUG=True
  23
  24 class TestNetworkAwareExtendHandshake(TestExtendHandshake):
  25     """
  26     Test our network awareness code that tries to detect if we're behind
  27     the same NAT and if so, connect via the internal network.
  28
  29     See BitTornado/BT1/Connecter.py
  30     """
  31
  32     def setUp(self):
  33         """ override TestAsServer """
  34         TestExtendHandshake.setUp(self)
  35
  36     def setUpPreSession(self):
  37         """ override TestAsServer """
  38         TestExtendHandshake.setUpPreSession(self)
  39
  40     def setUpPostSession(self):
  41         """ override TestAsServer """
  42         TestExtendHandshake.setUpPostSession(self)
  43
  44         # Create a fake "internal network interface"
  45         self.setUpMyListenSocket()
  46
  47         self.myid = "R------andomPeer4811"
  48
  49     def setUpDownloadConfig(self):
  50         dscfg = TestExtendHandshake.setUpDownloadConfig(self)
  51
  52         dscfg.set_same_nat_try_internal(True)
  53         dscfg.set_unchoke_bias_for_internal(481)
  54         return dscfg
  55
  56     def setUpMyListenSocket(self):
  57         self.destport = 4811
  58
  59         # Start our server side, to with Tribler will try to connect
  60         self.destss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  61         self.destss.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  62         self.destss.bind(('', self.destport))
  63         self.destss.listen(1)
  64
  65     def singtest_ext_ip_unknown(self):
  66         """ Send EXTEND hs to Tribler with yourip set to 127.0.0.1, so it
  67         appears we are using the same IP address. Tribler doesn't know its
  68         own external IP address, so it will do a loopback test to yourip
  69         to check. If this check succeeds it will connect to us via the
  70         internal network, i.e., our internal interface self.destss.
  71         """
  72         self.session.lm.upnp_ext_ip = None
  73         self.subtest_good_tribler_extend_hs()
  74
  75
  76     def singtest_ext_ip_known(self):
  77         """ Same as singtest_ext_ip_unknown() except no loopback test is needed
  78         as Tribler knows its external IP and it will be the same as the sent yourip.
  79         """
  80         self.session.lm.upnp_ext_ip = '127.0.0.1'
  81         self.subtest_good_tribler_extend_hs()
  82
  83
  84     #
  85     # Good EXTEND handshake message
  86     #
  87     def subtest_good_tribler_extend_hs(self):
  88         self._test_good(self.create_good_tribler_extend_hs,infohash=self.infohash)
  89
  90     def _test_good(self,msg_gen_func,options=None,infohash=None):
  91
  92         print >>sys.stderr,"test: test good, gen_func",msg_gen_func
  93
  94         if options is None and infohash is None:
  95             s = BTConnection('localhost',self.hisport,myid=self.myid)
  96         elif options is None:
  97             s = BTConnection('localhost',self.hisport,user_infohash=infohash,myid=self.myid)
  98         elif infohash is None:
  99             s = BTConnection('localhost',self.hisport,user_option_pattern=options,myid=self.myid)
 100         else:
 101             s = BTConnection('localhost',self.hisport,user_option_pattern=options,user_infohash=infohash,myid=self.myid)
 102         msg = msg_gen_func()
 103         s.send(msg)
 104         s.read_handshake_medium_rare()
 105         time.sleep(5)
 106
 107         # Tribler should send an EXTEND message back
 108         try:
 109             s.s.settimeout(10.0)
 110             resp = s.recv()
 111             self.assert_(len(resp) > 0)
 112             print >>sys.stderr,"test: Got reply",getMessageName(resp[0])
 113             self.assert_(resp[0] == EXTEND)
 114             self.check_tribler_extend_hs(resp[1:])
 115             #s.close()
 116         except socket.timeout:
 117             print >> sys.stderr,"test: Timeout, bad, peer didn't reply with EXTEND message"
 118             self.assert_(False)
 119
 120         # Tribler should try to connect to our internal interface
 121         self.destss.settimeout(10.0)
 122         conn, addr = self.destss.accept()
 123         s2 = BTConnection('',0,conn,user_infohash=self.infohash,myid=self.myid)
 124         s2.send(INTERESTED)
 125         s2.read_handshake_medium_rare()
 126
 127         # Is it him?
 128         self.assert_(s.hisid == s2.hisid)
 129
 130         # He should close original conn
 131         try:
 132             while True:
 133                 resp = s.recv()
 134                 if len(resp) > 0:
 135                     print >>sys.stderr,"test: Got data on internal conn",getMessageName(resp[0])
 136                 else:
 137                     break
 138         except socket.timeout:
 139             self.assert_(False)
 140
 141         self.assert_(True)
 142
 143
 144     def create_good_tribler_extend_hs(self):
 145         d = {}
 146         d['m'] = {'Tr_OVERLAYSWARM':253}
 147         d['p'] = self.mylistenport
 148         d['v'] = 'Tribler 4.5.0'
 149         d['yourip'] = compact_ip('127.0.0.1')
 150         d['ipv4'] = compact_ip('224.4.8.1')
 151         bd = bencode(d)
 152         return EXTEND+chr(0)+bd
 153
 154
 155     def check_tribler_extend_hs(self,data):
 156         self.assert_(data[0] == chr(0))
 157         d = bdecode(data[1:])
 158         self.assert_(type(d) == DictType)
 159         self.assert_('m' in d.keys())
 160         m = d['m']
 161         self.assert_(type(m) == DictType)
 162         self.assert_('Tr_OVERLAYSWARM' in m.keys())
 163         val = m['Tr_OVERLAYSWARM']
 164         self.assert_(type(val) == IntType)
 165         self.assert_(val == 253)
 166
 167         print >>sys.stderr,"test: Reply is",`d`
 168
 169 def test_suite():
 170     suite = unittest.TestSuite()
 171     # We should run the tests in a separate Python interpreter to prevent
 172     # problems with our singleton classes, e.g. PeerDB, etc.
 173     if len(sys.argv) != 2:
 174         print "Usage: python test_na_extend_hs.py <method name>"
 175     else:
 176         suite.addTest(TestNetworkAwareExtendHandshake(sys.argv[1]))
 177
 178     return suite
 179
 180 def main():
 181     unittest.main(defaultTest='test_suite',argv=[sys.argv[0]])
 182
 183 if __name__ == "__main__":
 184     main()